xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Smart Model loading (#31)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Ayush Chaurasia
2022-11-01 04:22:12 +05:30
committed by GitHub
parent 1054819a59
commit 92c60758dd
4 changed files with 80 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
ultralytics/yolo/engine/model.py
View File

@ -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)

16
ultralytics/yolo/engine/trainer.py
View File

@ -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

20
ultralytics/yolo/utils/modeling/__init__.py
View File

@ -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:

19
ultralytics/yolo/v8/classify/train.py
View File

@ -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()