From d76d7af56696d8ed242fd68eab203bf9008f0e27 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 2 Jan 2023 15:42:01 +0100 Subject: [PATCH] Simplify usage, remove YOLO.new() and YOLO.load() --- tests/check_flops.py | 2 +- tests/test_model.py | 18 +++++------ ultralytics/yolo/engine/exporter.py | 2 +- ultralytics/yolo/engine/model.py | 49 ++++++++++++----------------- ultralytics/yolo/utils/__init__.py | 4 +-- 5 files changed, 33 insertions(+), 42 deletions(-) diff --git a/tests/check_flops.py b/tests/check_flops.py index 5ddebec..fd2bbd3 100644 --- a/tests/check_flops.py +++ b/tests/check_flops.py @@ -4,7 +4,7 @@ from ultralytics.yolo.utils import ROOT if __name__ == "__main__": for m in list((ROOT / 'yolo/v8/models').rglob('*.yaml')): try: - YOLO.new(m.name, verbose=True) + YOLO(m.name, verbose=True) except Exception as e: print(f'ERROR for {m}: {e}') diff --git a/tests/test_model.py b/tests/test_model.py index 119f32d..3277bbd 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -4,7 +4,7 @@ from ultralytics import YOLO def test_model_init(): - model = YOLO.new("yolov8n.yaml") + model = YOLO("yolov8n.yaml") model.info() try: YOLO() @@ -14,38 +14,38 @@ def test_model_init(): def test_model_forward(): - model = YOLO.new("yolov8n.yaml") + model = YOLO("yolov8n.yaml") img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512) model.forward(img) model(img) def test_model_info(): - model = YOLO.new("yolov8n.yaml") + model = YOLO("yolov8n.yaml") model.info() model = model.load("best.pt") model.info(verbose=True) def test_model_fuse(): - model = YOLO.new("yolov8n.yaml") + model = YOLO("yolov8n.yaml") model.fuse() model.load("best.pt") model.fuse() def test_visualize_preds(): - model = YOLO.load("best.pt") + model = YOLO("best.pt") model.predict(source="ultralytics/assets") def test_val(): - model = YOLO.load("best.pt") + model = YOLO("best.pt") model.val(data="coco128.yaml", imgsz=32) def test_model_resume(): - model = YOLO.new("yolov8n.yaml") + model = YOLO("yolov8n.yaml") model.train(epochs=1, imgsz=32, data="coco128.yaml") try: model.resume(task="detect") @@ -54,7 +54,7 @@ def test_model_resume(): def test_model_train_pretrained(): - model = YOLO.load("best.pt") + model = YOLO("best.pt") model.train(data="coco128.yaml", epochs=1, imgsz=32) model = model.new("yolov8n.yaml") model.train(data="coco128.yaml", epochs=1, imgsz=32) @@ -83,7 +83,7 @@ def test_exports(): print(export_formats()) - model = YOLO.new("yolov8n.yaml") + model = YOLO("yolov8n.yaml") model.export(format='torchscript') model.export(format='onnx') model.export(format='openvino') diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py index 71e9cb8..e518841 100644 --- a/ultralytics/yolo/engine/exporter.py +++ b/ultralytics/yolo/engine/exporter.py @@ -23,7 +23,7 @@ Requirements: Python: from ultralytics import YOLO - model = YOLO.new('yolov8n.yaml') + model = YOLO('yolov8n.yaml') results = model.export(format='onnx') CLI: diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index da53a62..20b9514 100644 --- a/ultralytics/yolo/engine/model.py +++ b/ultralytics/yolo/engine/model.py @@ -1,4 +1,5 @@ import torch +from pathlib import Path from ultralytics import yolo # noqa from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights @@ -27,19 +28,15 @@ class YOLO: A python interface which emulates a model-like behaviour by wrapping trainers. """ - __init_key = object() # used to ensure proper initialization - def __init__(self, init_key=None, type="v8") -> None: + def __init__(self, model='yolov8n.yaml', type="v8") -> None: """ Initializes the YOLO object. Args: - init_key (object): used to ensure proper initialization. Defaults to None. + model (str, Path): model to load or create type (str): Type/version of models to use. Defaults to "v8". """ - if init_key != YOLO.__init_key: - raise SyntaxError(HELP_MSG) - self.type = type self.ModelClass = None # model class self.TrainerClass = None # trainer class @@ -53,8 +50,10 @@ class YOLO: self.overrides = {} # overrides for trainer object self.init_disabled = False # disable model initialization - @classmethod - def new(cls, cfg: str, verbose=True): + # Load or create new YOLO model + {'.pt': self._load, '.yaml': self._new}[Path(model).suffix](model) + + def _new(self, cfg: str, verbose=True): """ Initializes a new model and infers the task type from the model definitions. @@ -64,34 +63,26 @@ class YOLO: """ cfg = check_yaml(cfg) # check YAML cfg_dict = yaml_load(cfg) # model dict - obj = cls(init_key=cls.__init_key) - obj.task = guess_task_from_head(cfg_dict["head"][-1][-2]) - obj.ModelClass, obj.TrainerClass, obj.ValidatorClass, obj.PredictorClass = obj._guess_ops_from_task(obj.task) - obj.model = obj.ModelClass(cfg_dict, verbose=verbose) # initialize - obj.cfg = cfg - - return obj + self.task = guess_task_from_head(cfg_dict["head"][-1][-2]) + self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = \ + self._guess_ops_from_task(self.task) + self.model = self.ModelClass(cfg_dict, verbose=verbose) # initialize + self.cfg = cfg - @classmethod - def load(cls, weights: str): + 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 """ - obj = cls(init_key=cls.__init_key) - obj.ckpt = torch.load(weights, map_location="cpu") - obj.task = obj.ckpt["train_args"]["task"] - obj.overrides = dict(obj.ckpt["train_args"]) - obj.overrides["device"] = '' # reset device - LOGGER.info("Device has been reset to ''") - - obj.ModelClass, obj.TrainerClass, obj.ValidatorClass, obj.PredictorClass = obj._guess_ops_from_task( - task=obj.task) - obj.model = attempt_load_weights(weights) - - return obj + self.ckpt = torch.load(weights, map_location="cpu") + self.task = self.ckpt["train_args"]["task"] + self.overrides = dict(self.ckpt["train_args"]) + self.overrides["device"] = '' # reset device + self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = \ + self._guess_ops_from_task(self.task) + self.model = attempt_load_weights(weights, fuse=False) def reset(self): """ diff --git a/ultralytics/yolo/utils/__init__.py b/ultralytics/yolo/utils/__init__.py index 9e6fd3e..856358c 100644 --- a/ultralytics/yolo/utils/__init__.py +++ b/ultralytics/yolo/utils/__init__.py @@ -36,8 +36,8 @@ HELP_MSG = \ from ultralytics import YOLO - model = YOLO.new('yolov8n.yaml') # create a new model from scratch - model = YOLO.load('yolov8n.pt') # load a pretrained model (recommended for best training results) + model = YOLO('yolov8n.yaml') # build a new model from scratch + model = YOLO('yolov8n.pt') # load a pretrained model (recommended for best training results) results = model.train(data='coco128.yaml') # train the model results = model.val() # evaluate model performance on the validation set results = model.predict(source='bus.jpg') # predict on an image