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

ultralytics/yolo/utils/modeling/__init__.py

@@ -1,10 +1,10 @@
 import contextlib
 
+import torchvision
 import yaml
 
 from ultralytics.yolo.utils.downloads import attempt_download
-
-from .modules import *
+from ultralytics.yolo.utils.modeling.modules import *
 
 
 def attempt_load_weights(weights, device=None, inplace=True, fuse=True):
@@ -26,7 +26,7 @@ def attempt_load_weights(weights, device=None, inplace=True, fuse=True):
     # Module compatibility updates
     for m in model.modules():
         t = type(m)
-        if t in (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model):
+        if t in (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect):
             m.inplace = inplace  # torch 1.7.0 compatibility
             if t is Detect and not isinstance(m.anchor_grid, list):
                 delattr(m, 'anchor_grid')
@@ -107,6 +107,20 @@ def parse_model(d, ch):  # model_dict, input_channels(3)
     return nn.Sequential(*layers), sorted(save)
 
 
+def get_model(model: str):
+    if model.endswith(".pt"):
+        model = model.split(".")[0]
+
+    if Path(model + ".pt").is_file():
+        trained_model = torch.load(model + ".pt", map_location='cpu')
+    elif model in torchvision.models.__dict__:  # try torch hub classifier models
+        trained_model = torch.hub.load("pytorch/vision", model, pretrained=True)
+    else:
+        model_ckpt = attempt_download(model + ".pt")  # try ultralytics assets
+        trained_model = torch.load(model_ckpt, map_location='cpu')
+    return trained_model
+
+
 def yaml_load(file='data.yaml'):
     # Single-line safe yaml loading
     with open(file, errors='ignore') as f:

ultralytics/yolo/v8/classify/train.py

@@ -41,21 +41,6 @@ class ClassificationTrainer(BaseTrainer):
     def get_dataloader(self, dataset_path, batch_size=None, rank=0):
         return build_classification_dataloader(path=dataset_path, batch_size=self.args.batch_size, rank=rank)
 
-    def get_model(self, model, pretrained):
-        # temp. minimal. only supports torchvision models
-        model = self.args.model
-        if model in torchvision.models.__dict__:  # TorchVision models i.e. resnet50, efficientnet_b0
-            model = torchvision.models.__dict__[model](weights='IMAGENET1K_V1' if pretrained else None)
-        else:
-            raise ModuleNotFoundError(f'--model {model} not found.')
-        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  # for training
-
-        return model
-
     def get_validator(self):
         return v8.classify.ClassificationValidator(self.test_loader, self.device, logger=self.console)
 
@@ -65,8 +50,8 @@ class ClassificationTrainer(BaseTrainer):
 
 @hydra.main(version_base=None, config_path=CONFIG_PATH_ABS, config_name=str(DEFAULT_CONFIG).split(".")[0])
 def train(cfg):
-    cfg.model = cfg.model or "squeezenet1_0"
-    cfg.data = cfg.data or "imagenette"  # or yolo.ClassificationDataset("mnist")
+    cfg.model = cfg.model or "resnet18"
+    cfg.data = cfg.data or "imagenette160"  # or yolo.ClassificationDataset("mnist")
     trainer = ClassificationTrainer(cfg)
     trainer.train()