From 92c60758ddfe3cf8860a9caa4068892a76ed766a Mon Sep 17 00:00:00 2001 From: Ayush Chaurasia Date: Tue, 1 Nov 2022 04:22:12 +0530 Subject: [PATCH] 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 | 67 +++++++++++++++------ ultralytics/yolo/engine/trainer.py | 16 ++++- ultralytics/yolo/utils/modeling/__init__.py | 20 +++++- ultralytics/yolo/v8/classify/train.py | 19 +----- 4 files changed, 80 insertions(+), 42 deletions(-) diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 838014b77..d5e5ca4d8 100644 --- a/ultralytics/yolo/engine/model.py +++ b/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_class, trainer_class + - return model(cfg), trainer +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) diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index 875bc351d..44decbd49 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/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 diff --git a/ultralytics/yolo/utils/modeling/__init__.py b/ultralytics/yolo/utils/modeling/__init__.py index adc1f5f34..c3328c7b3 100644 --- a/ultralytics/yolo/utils/modeling/__init__.py +++ b/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: diff --git a/ultralytics/yolo/v8/classify/train.py b/ultralytics/yolo/v8/classify/train.py index 534dbfd32..6cc21bbe3 100644 --- a/ultralytics/yolo/v8/classify/train.py +++ b/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()