xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update docs (#71)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Ayush Chaurasia
2022-12-12 09:21:00 +05:30
committed by GitHub
parent e629335f6d
commit d85b44f259
11 changed files with 286 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
ultralytics/yolo/engine/model.py
View File

@ -1,26 +1,31 @@
"""
Top-level YOLO model interface. First principle usage example - https://github.com/ultralytics/ultralytics/issues/13
"""
import torch
import yaml
from ultralytics import yolo
from ultralytics.yolo.utils import LOGGER
from ultralytics.yolo.utils.checks import check_yaml
from ultralytics.yolo.utils.files import yaml_load
from ultralytics.yolo.utils.modeling import attempt_load_weights
from ultralytics.yolo.utils.modeling.tasks import ClassificationModel, DetectionModel, SegmentationModel
# map head: [model, trainer]
MODEL_MAP = {
"classify": [ClassificationModel, 'yolo.VERSION.classify.ClassificationTrainer'],
"detect": [DetectionModel, 'yolo.VERSION.detect.DetectionTrainer'],
"segment": [SegmentationModel, 'yolo.VERSION.segment.SegmentationTrainer']}
"classify": [ClassificationModel, 'yolo.TYPE.classify.ClassificationTrainer'],
"detect": [DetectionModel, 'yolo.TYPE.detect.DetectionTrainer'],
"segment": [SegmentationModel, 'yolo.TYPE.segment.SegmentationTrainer']}
class YOLO:
"""
Python interface which emulates a model-like behaviour by wrapping trainers.
"""
def __init__(self, version=8) -> None:
self.version = version
def __init__(self, type="v8") -> None:
"""
Args:
type (str): Type/version of models to use
"""
self.type = type
self.ModelClass = None
self.TrainerClass = None
self.model = None
@ -29,20 +34,36 @@ class YOLO:
self.ckpt = None
def new(self, cfg: str):
"""
Initializes a new model and infers the task type from the model definitions
Args:
cfg (str): model configuration file
"""
cfg = check_yaml(cfg) # check YAML
with open(cfg, encoding='ascii', errors='ignore') as f:
cfg = yaml.safe_load(f) # model dict
self.ModelClass, self.TrainerClass, self.task = self._guess_model_trainer_and_task(cfg["head"][-1][-2])
self.model = self.ModelClass(cfg) # initialize
def load(self, weights):
def load(self, weights: str):
"""
Initializes a new model and infers the task type from the model head
Args:
weights (str): model checkpoint to be loaded
"""
self.ckpt = torch.load(weights, map_location="cpu")
self.task = self.ckpt["train_args"]["task"]
_, trainer_class_literal = MODEL_MAP[self.task]
self.TrainerClass = eval(trainer_class_literal.replace("VERSION", f"v{self.version}"))
self.TrainerClass = eval(trainer_class_literal.replace("TYPE", f"v{self.type}"))
self.model = attempt_load_weights(weights)
def reset(self):
"""
Resets the model modules .
"""
for m in self.model.modules():
if hasattr(m, 'reset_parameters'):
m.reset_parameters()
@ -50,32 +71,46 @@ class YOLO:
p.requires_grad = True
def train(self, **kwargs):
if 'data' not in kwargs:
raise Exception("data is required to train")
"""
Trains the model on given dataset.
Args:
**kwargs (Any): Any number of arguments representing the training configuration. List of all args can be found in 'config' section.
You can pass all arguments as a yaml file in `cfg`. Other args are ignored if `cfg` file is passed
"""
if not self.model and not self.ckpt:
raise Exception("model not initialized. Use .new() or .load()")
kwargs["task"] = self.task
kwargs["mode"] = "train"
self.trainer = self.TrainerClass(overrides=kwargs)
overrides = kwargs
if kwargs.get("cfg"):
LOGGER.info(f"cfg file passed. Overriding default params with {kwargs['cfg']}.")
overrides = yaml_load(check_yaml(kwargs["cfg"]))
overrides["task"] = self.task
overrides["mode"] = "train"
if not overrides.get("data"):
raise Exception("dataset not provided! Please check if you have defined `data` in you configs")
self.trainer = self.TrainerClass(overrides=overrides)
# load pre-trained weights if found, else use the loaded model
self.trainer.model = self.trainer.load_model(weights=self.ckpt) if self.ckpt else self.model
self.trainer.train()
def resume(self, task=None, model=None):
if not task:
raise Exception(
"pass the task type and/or model(optional) from which you want to resume: `model.resume(task="
")`")
def resume(self, task, model=None):
"""
Resume a training task.
Args:
task (str): The task type you want to resume. Automatically finds the last run to resume if `model` is not specified.
model (str): [Optional] The model checkpoint to resume from. If not found, the last run of the given task type is resumed.
"""
if task.lower() not in MODEL_MAP:
raise Exception(f"unrecognised task - {task}. Supported tasks are {MODEL_MAP.keys()}")
_, trainer_class_literal = MODEL_MAP[task.lower()]
self.TrainerClass = eval(trainer_class_literal.replace("VERSION", f"v{self.version}"))
self.TrainerClass = eval(trainer_class_literal.replace("TYPE", f"v{self.type}"))
self.trainer = self.TrainerClass(overrides={"task": task.lower(), "resume": model if model else True})
self.trainer.train()
def _guess_model_trainer_and_task(self, head):
# TODO: warn
task = None
if head.lower() in ["classify", "classifier", "cls", "fc"]:
task = "classify"
@ -85,7 +120,7 @@ class YOLO:
task = "segment"
model_class, trainer_class = MODEL_MAP[task]
# warning: eval is unsafe. Use with caution
trainer_class = eval(trainer_class.replace("VERSION", f"v{self.version}"))
trainer_class = eval(trainer_class.replace("TYPE", f"{self.type}"))
return model_class, trainer_class, task

4
ultralytics/yolo/engine/trainer.py
View File

@ -35,8 +35,8 @@ RANK = int(os.getenv('RANK', -1))
class BaseTrainer:
def __init__(self, config=DEFAULT_CONFIG, overrides={}):
self.args = get_config(config, overrides)
def __init__(self, cfg=DEFAULT_CONFIG, overrides={}):
self.args = get_config(cfg, overrides)
self.check_resume()
init_seeds(self.args.seed + 1 + RANK, deterministic=self.args.deterministic)