diff --git a/ultralytics/yolo/__init__.py b/ultralytics/yolo/__init__.py index 886a846..37f36e0 100644 --- a/ultralytics/yolo/__init__.py +++ b/ultralytics/yolo/__init__.py @@ -1,3 +1,5 @@ +from ultralytics.yolo import v8 + from .engine.model import YOLO from .engine.trainer import BaseTrainer from .engine.validator import BaseValidator diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 3a70a2d..163210d 100644 --- a/ultralytics/yolo/engine/model.py +++ b/ultralytics/yolo/engine/model.py @@ -1,55 +1,45 @@ """ Top-level YOLO model interface. First principle usage example - https://github.com/ultralytics/ultralytics/issues/13 """ +import torch import yaml from ultralytics.yolo.utils import LOGGER from ultralytics.yolo.utils.checks import check_yaml -from ultralytics.yolo.utils.modeling import get_model +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.train.ClassificationTrainer'], - "detect": [ClassificationModel, 'yolo.VERSION.classify.train.ClassificationTrainer'], # temp - "segment": []} + "classify": [ClassificationModel, 'yolo.VERSION.classify.ClassificationTrainer'], + "detect": [DetectionModel, 'yolo.VERSION.detect.DetectionTrainer'], + "segment": [SegmentationModel, 'yolo.VERSION.segment.SegmentationTrainer']} class YOLO: - def __init__(self, task=None, version=8) -> None: + def __init__(self, version=8) -> None: self.version = version self.ModelClass = None self.TrainerClass = None self.model = None - self.pretrained_weights = None - if task: - if task.lower() not in MODEL_MAP: - raise Exception(f"Unsupported task {task}. The supported tasks are: \n {MODEL_MAP.keys()}") - self.ModelClass, self.TrainerClass = MODEL_MAP[task] - self.TrainerClass = eval(self.trainer.replace("VERSION", f"v{self.version}")) + self.trainer = None + self.task = None + self.ckpt = None def new(self, cfg: str): cfg = check_yaml(cfg) # check YAML - if self.model: - self.model = self.model(cfg) - else: - with open(cfg, encoding='ascii', errors='ignore') as f: - cfg = yaml.safe_load(f) # model dict - self.ModelClass, self.TrainerClass = self._get_model_and_trainer(cfg["head"]) - self.model = self.ModelClass(cfg) # initialize + 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, autodownload=True): - if not isinstance(self.pretrained_weights, type(None)): - LOGGER.info("Overwriting weights") - # TODO: weights = smart_file_loader(weights) - if self.model: - self.model.load(weights) - LOGGER.info("Checkpoint loaded successfully") - else: - self.model = get_model(weights) - self.ModelClass, self.TrainerClass = self._guess_model_and_trainer(list(self.model.named_children())) - self.pretrained_weights = weights + def load(self, weights): + 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.model = attempt_load_weights(weights) def reset(self): for m in self.model.modules(): @@ -61,16 +51,31 @@ class YOLO: def train(self, **kwargs): if 'data' not in kwargs: raise Exception("data is required to train") - if not self.model: + if not self.model and not self.ckpt: raise Exception("model not initialized. Use .new() or .load()") - # kwargs["model"] = self.model - trainer = self.TrainerClass(overrides=kwargs) - trainer.model = self.model - trainer.train() - def _guess_model_and_trainer(self, cfg): + kwargs["task"] = self.task + kwargs["mode"] = "train" + self.trainer = self.TrainerClass(overrides=kwargs) + # 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=" + ")`") + 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.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 - head = cfg[-1][-2] + task = None if head.lower() in ["classify", "classifier", "cls", "fc"]: task = "classify" if head.lower() in ["detect"]: @@ -81,11 +86,9 @@ class YOLO: # warning: eval is unsafe. Use with caution trainer_class = eval(trainer_class.replace("VERSION", f"v{self.version}")) - return model_class, trainer_class + return model_class, trainer_class, task - -if __name__ == "__main__": - model = YOLO() - # model.new("assets/dummy_model.yaml") - model.load("yolov5n-cls.pt") - model.train(data="imagenette160", epochs=1, lr0=0.01) + def __call__(self, imgs): + if not self.model: + LOGGER.info("model not initialized!") + return self.model(imgs) diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index 97bcf1c..4dcbf4c 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/ultralytics/yolo/engine/trainer.py @@ -8,7 +8,6 @@ from collections import defaultdict from copy import deepcopy from datetime import datetime from pathlib import Path -from typing import Dict, Union import numpy as np import torch @@ -28,7 +27,6 @@ from ultralytics.yolo.utils import LOGGER, ROOT, TQDM_BAR_FORMAT, colorstr from ultralytics.yolo.utils.checks import check_file, print_args from ultralytics.yolo.utils.configs import get_config from ultralytics.yolo.utils.files import get_latest_run, increment_path, save_yaml -from ultralytics.yolo.utils.modeling import get_model from ultralytics.yolo.utils.torch_utils import ModelEMA, de_parallel, init_seeds, one_cycle, strip_optimizer DEFAULT_CONFIG = ROOT / "yolo/utils/configs/default.yaml" @@ -63,6 +61,7 @@ class BaseTrainer: self.scaler = amp.GradScaler(enabled=self.device.type != 'cpu') # Model and Dataloaders. + self.model = self.args.model self.data = self.args.data if self.data.endswith(".yaml"): self.data = check_dataset_yaml(self.data) @@ -125,6 +124,7 @@ class BaseTrainer: """ # model ckpt = self.setup_model() + self.model = self.model.to(self.device) self.set_model_attributes() if world_size > 1: self.model = DDP(self.model, device_ids=[rank]) @@ -288,13 +288,16 @@ class BaseTrainer: """ load/create/download model for any task """ - model = self.args.model + if isinstance(self.model, torch.nn.Module): # if loaded model is passed + return + # We should improve the code flow here. This function looks hacky + model = self.model pretrained = not (str(model).endswith(".yaml")) # config if not pretrained: model = check_file(model) ckpt = self.load_ckpt(model) if pretrained else None - self.model = self.load_model(model_cfg=None if pretrained else model, weights=ckpt).to(self.device) # model + self.model = self.load_model(model_cfg=None if pretrained else model, weights=ckpt) # model return ckpt def load_ckpt(self, ckpt): @@ -402,7 +405,7 @@ class BaseTrainer: last = Path(check_file(resume) if isinstance(resume, str) else get_latest_run()) args_yaml = last.parent.parent / 'args.yaml' # train options yaml if args_yaml.is_file(): - args = self._get_config(args_yaml) # replace + args = get_config(args_yaml) # replace args.model, args.resume, args.exist_ok = str(last), True, True # reinstate self.args = args @@ -424,8 +427,7 @@ class BaseTrainer: f'Resuming training from {self.args.model} from epoch {start_epoch} to {self.epochs} total epochs') if self.epochs < start_epoch: LOGGER.info( - f"{self.args.model} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {self.epochs} more epochs." - ) + f"{self.model} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {self.epochs} more epochs.") self.epochs += ckpt['epoch'] # finetune additional epochs self.best_fitness = best_fitness self.start_epoch = start_epoch @@ -460,9 +462,3 @@ def build_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5): LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}) with parameter groups " f"{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias") return optimizer - - -# Dummy validator -def val(trainer: BaseTrainer): - trainer.console.info("validating") - return {"metric_1": 0.1, "metric_2": 0.2, "fitness": 1} diff --git a/ultralytics/yolo/v8/classify/train.py b/ultralytics/yolo/v8/classify/train.py index 813278d..95fe4ef 100644 --- a/ultralytics/yolo/v8/classify/train.py +++ b/ultralytics/yolo/v8/classify/train.py @@ -13,8 +13,10 @@ class ClassificationTrainer(BaseTrainer): def set_model_attributes(self): self.model.names = self.data["names"] - def load_model(self, model_cfg, weights): + def load_model(self, model_cfg=None, weights=None): # TODO: why treat clf models as unique. We should have clf yamls? + if isinstance(weights, dict): # yolo ckpt + weights = weights["model"] if weights and not weights.__class__.__name__.startswith("yolo"): # torchvision model = weights else: diff --git a/ultralytics/yolo/v8/detect/train.py b/ultralytics/yolo/v8/detect/train.py index 8021b78..553f0d9 100644 --- a/ultralytics/yolo/v8/detect/train.py +++ b/ultralytics/yolo/v8/detect/train.py @@ -15,7 +15,7 @@ from .val import DetectionValidator # BaseTrainer python usage class DetectionTrainer(SegmentationTrainer): - def load_model(self, model_cfg, weights): + def load_model(self, model_cfg=None, weights=None): model = DetectionModel(model_cfg or weights["model"].yaml, ch=3, nc=self.data["nc"], diff --git a/ultralytics/yolo/v8/segment/train.py b/ultralytics/yolo/v8/segment/train.py index 95bc417..0286ce1 100644 --- a/ultralytics/yolo/v8/segment/train.py +++ b/ultralytics/yolo/v8/segment/train.py @@ -26,7 +26,7 @@ class SegmentationTrainer(BaseTrainer): batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255 return batch - def load_model(self, model_cfg, weights): + def load_model(self, model_cfg=None, weights=None): model = SegmentationModel(model_cfg or weights["model"].yaml, ch=3, nc=self.data["nc"],