Smart Model loading (#31)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/yolo/engine/model.py

@@ -1,32 +1,44 @@
 """
 Top-level YOLO model interface. First principle usage example - https://github.com/ultralytics/ultralytics/issues/13
 """
-import torch
 import yaml
 
 import ultralytics.yolo as yolo
 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.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.train.ClassificationTrainer'],
+    "detect": [ClassificationModel, 'yolo.VERSION.classify.train.ClassificationTrainer'],  # temp
+    "segment": []}
 
 
 class YOLO:
 
-    def __init__(self, version=8) -> None:
+    def __init__(self, task=None, version=8) -> None:
         self.version = version
+        self.ModelClass = None
+        self.TrainerClass = None
         self.model = None
-        self.trainer = 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}"))
 
     def new(self, cfg: str):
         cfg = check_yaml(cfg)  # check YAML
-        self.model, self.trainer = self._get_model_and_trainer(cfg)
+        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
 
     def load(self, weights, autodownload=True):
         if not isinstance(self.pretrained_weights, type(None)):
@@ -36,28 +48,45 @@ class YOLO:
             self.model.load(weights)
             LOGGER.info("Checkpoint loaded successfully")
         else:
-            # TODO: infer model and trainer
-            pass
-
+            self.model = get_model(weights)
+            self.ModelClass, self.TrainerClass = self._guess_model_and_trainer(list(self.model.named_children()))
         self.pretrained_weights = weights
 
     def reset(self):
-        pass
+        for m in self.model.modules():
+            if hasattr(m, 'reset_parameters'):
+                m.reset_parameters()
+        for p in self.model.parameters():
+            p.requires_grad = True
 
     def train(self, **kwargs):
         if 'data' not in kwargs:
             raise Exception("data is required to train")
         if not self.model:
             raise Exception("model not initialized. Use .new() or .load()")
-        kwargs["model"] = self.model
-        trainer = self.trainer(overrides=kwargs)
+        # kwargs["model"] = self.model
+        trainer = self.TrainerClass(overrides=kwargs)
+        trainer.model = self.model
         trainer.train()
 
-    def _get_model_and_trainer(self, cfg):
-        with open(cfg, encoding='ascii', errors='ignore') as f:
-            cfg = yaml.safe_load(f)  # model dict
-        model, trainer = MODEL_MAP[cfg["head"][-1][-2]]
+    def _guess_model_and_trainer(self, cfg):
+        # TODO: warn
+        head = cfg[-1][-2]
+        if head.lower() in ["classify", "classifier", "cls", "fc"]:
+            task = "classify"
+        if head.lower() in ["detect"]:
+            task = "detect"
+        if head.lower() in ["segment"]:
+            task = "segment"
+        model_class, trainer_class = MODEL_MAP[task]
         # warning: eval is unsafe. Use with caution
-        trainer = eval(trainer.replace("VERSION", f"v{self.version}"))
+        trainer_class = eval(trainer_class.replace("VERSION", f"v{self.version}"))
 
-        return model(cfg), trainer
+        return model_class, trainer_class
+
+
+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)

ultralytics/yolo/engine/trainer.py

@@ -22,6 +22,7 @@ import ultralytics.yolo.utils as utils
 import ultralytics.yolo.utils.loggers as loggers
 from ultralytics.yolo.utils import LOGGER, ROOT
 from ultralytics.yolo.utils.files import increment_path, save_yaml
+from ultralytics.yolo.utils.modeling import get_model
 
 CONFIG_PATH_ABS = ROOT / "yolo/utils/configs"
 DEFAULT_CONFIG = "defaults.yaml"
@@ -33,6 +34,7 @@ class BaseTrainer:
         self.console = LOGGER
         self.args = self._get_config(config, overrides)
         self.validator = None
+        self.model = None
         self.callbacks = defaultdict(list)
         self.console.info(f"Training config: \n args: \n {self.args}")  # to debug
         # Directories
@@ -51,7 +53,8 @@ class BaseTrainer:
 
         # Model and Dataloaders.
         self.trainset, self.testset = self.get_dataset(self.args.data)
-        self.model = self.get_model(self.args.model, self.args.pretrained).to(self.device)
+        if self.args.model is not None:
+            self.model = self.get_model(self.args.model, self.args.pretrained).to(self.device)
 
         # epoch level metrics
         self.metrics = {}  # handle metrics returned by validator
@@ -225,11 +228,18 @@ class BaseTrainer:
         """
         pass
 
-    def get_model(self, model, pretrained=True):
+    def get_model(self, model, pretrained):
         """
         load/create/download model for any task
         """
-        pass
+        model = get_model(model)
+        for m in model.modules():
+            if not pretrained and hasattr(m, 'reset_parameters'):
+                m.reset_parameters()
+        for p in model.parameters():
+            p.requires_grad = True
+
+        return model
 
     def get_validator(self):
         pass