mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 05:24:22 +08:00
Smart Model loading (#31)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
parent
1054819a59
commit
92c60758dd
@ -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)
|
||||
|
@ -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
|
||||
|
@ -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:
|
||||
|
@ -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()
|
||||
|
||||
|
Loading…
x
Reference in New Issue
Block a user