From c3d961fb0340b1c2c5c64df684e338ece821c009 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 2 Jan 2023 17:37:23 +0100 Subject: [PATCH] Unified model loading with backwards compatibility (#132) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/quickstart.md | 8 ++-- tests/test_model.py | 58 ++++++++++++++++------------ ultralytics/hub/utils.py | 2 +- ultralytics/nn/tasks.py | 18 +++++---- ultralytics/yolo/engine/exporter.py | 6 +-- ultralytics/yolo/engine/model.py | 9 ++--- ultralytics/yolo/engine/validator.py | 3 ++ ultralytics/yolo/utils/__init__.py | 5 ++- ultralytics/yolo/v8/detect/train.py | 2 +- ultralytics/yolo/v8/segment/train.py | 2 +- 10 files changed, 64 insertions(+), 49 deletions(-) diff --git a/docs/quickstart.md b/docs/quickstart.md index d73ac8a..cf796c7 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -42,9 +42,9 @@ Ultralytics YOLO comes with pythonic Model and Trainer interface. import ultralytics from ultralytics import YOLO - model = YOLO("s-seg.yaml") # automatically detects task type - model = YOLO("s-seg.pt") # load checkpoint - model.train(data="coco128-segments", epochs=1, lr0=0.01, ...) - model.train(data="coco128-segments", epochs=1, lr0=0.01, device="0,1,2,3") # DDP mode + model = YOLO("yolov8n-seg.yaml") # automatically detects task type + model = YOLO("yolov8n.pt") # load checkpoint + model.train(data="coco128-seg.yaml", epochs=1, lr0=0.01, ...) + model.train(data="coco128-seg.yaml", epochs=1, lr0=0.01, device="0,1,2,3") # DDP mode ``` [API Guide](sdk.md){ .md-button .md-button--primary} diff --git a/tests/test_model.py b/tests/test_model.py index 7089993..989d296 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -1,11 +1,7 @@ import torch from ultralytics import YOLO - - -def test_model_init(): - model = YOLO("yolov8n.yaml") - model.info() +from ultralytics.yolo.utils import ROOT def test_model_forward(): @@ -29,9 +25,9 @@ def test_model_fuse(): model.fuse() -def test_visualize_preds(): +def test_predict_dir(): model = YOLO("yolov8n.pt") - model.predict(source="ultralytics/assets") + model.predict(source=ROOT / "assets") def test_val(): @@ -39,7 +35,7 @@ def test_val(): model.val(data="coco128.yaml", imgsz=32) -def test_model_resume(): +def test_train_resume(): model = YOLO("yolov8n.yaml") model.train(epochs=1, imgsz=32, data="coco128.yaml") try: @@ -48,16 +44,21 @@ def test_model_resume(): print("Successfully caught resume assert!") -def test_model_train_pretrained(): - model = YOLO("yolov8n.pt") - model.train(data="coco128.yaml", epochs=1, imgsz=32) +def test_train_scratch(): model = YOLO("yolov8n.yaml") model.train(data="coco128.yaml", epochs=1, imgsz=32) img = torch.rand(1, 3, 320, 320) model(img) -def test_exports(): +def test_train_pretrained(): + model = YOLO("yolov8n.pt") + model.train(data="coco128.yaml", epochs=1, imgsz=32) + img = torch.rand(1, 3, 320, 320) + model(img) + + +def test_export_torchscript(): """ Format Argument Suffix CPU GPU 0 PyTorch - .pt True True @@ -74,26 +75,35 @@ def test_exports(): 11 PaddlePaddle paddle _paddle_model True True """ from ultralytics.yolo.engine.exporter import export_formats - print(export_formats()) model = YOLO("yolov8n.yaml") model.export(format='torchscript') + + +def test_export_onnx(): + model = YOLO("yolov8n.yaml") model.export(format='onnx') + + +def test_export_openvino(): + model = YOLO("yolov8n.yaml") model.export(format='openvino') + + +def test_export_coreml(): + model = YOLO("yolov8n.yaml") model.export(format='coreml') - model.export(format='paddle') -def test(): - test_model_forward() - test_model_info() - test_model_fuse() - test_visualize_preds() - test_val() - test_model_resume() - test_model_train_pretrained() +def test_export_paddle(): + model = YOLO("yolov8n.yaml") + model.export(format='paddle') -if __name__ == "__main__": - test() +# def run_all_tests(): # do not name function test_... +# pass +# +# +# if __name__ == "__main__": +# run_all_tests() diff --git a/ultralytics/hub/utils.py b/ultralytics/hub/utils.py index 348d0a4..34c06f4 100644 --- a/ultralytics/hub/utils.py +++ b/ultralytics/hub/utils.py @@ -124,7 +124,7 @@ def smart_request(*args, retry=3, timeout=30, thread=True, code=-1, method="post return func(*args, **kwargs) -def sync_analytics(cfg, all_keys=False, enabled=True): +def sync_analytics(cfg, all_keys=False, enabled=False): """ Sync analytics data if enabled in the global settings diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index afe0bfa..c1f6b9b 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -10,11 +10,13 @@ import torchvision from ultralytics.nn.modules import (C1, C2, C3, C3TR, SPP, SPPF, Bottleneck, BottleneckCSP, C2f, C3Ghost, C3x, Classify, Concat, Conv, ConvTranspose, Detect, DWConv, DWConvTranspose2d, Ensemble, Focus, GhostBottleneck, GhostConv, Segment) -from ultralytics.yolo.utils import LOGGER, colorstr, yaml_load +from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, colorstr, yaml_load from ultralytics.yolo.utils.checks import check_yaml from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, initialize_weights, intersect_dicts, make_divisible, model_info, scale_img, time_sync) +DEFAULT_CONFIG_DICT = yaml_load(DEFAULT_CONFIG, append_filename=False) + class BaseModel(nn.Module): ''' @@ -211,7 +213,7 @@ class DetectionModel(BaseModel): return y def load(self, weights, verbose=True): - csd = weights['model'].float().state_dict() # checkpoint state_dict as FP32 + csd = weights.float().state_dict() # checkpoint state_dict as FP32 csd = intersect_dicts(csd, self.state_dict()) # intersect self.load_state_dict(csd, strict=False) # load if verbose: @@ -281,21 +283,21 @@ class ClassificationModel(BaseModel): # Functions ------------------------------------------------------------------------------------------------------------ -def attempt_load_weights(weights, device=None, inplace=True, fuse=True): +def attempt_load_weights(weights, device=None, inplace=True, fuse=False): # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a from ultralytics.yolo.utils.downloads import attempt_download + default_keys = DEFAULT_CONFIG_DICT.keys() model = Ensemble() for w in weights if isinstance(weights, list) else [weights]: ckpt = torch.load(attempt_download(w), map_location='cpu') # load + args = {**DEFAULT_CONFIG_DICT, **ckpt['train_args']} ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model # Model compatibility updates - if not hasattr(ckpt, 'stride'): - ckpt.stride = torch.tensor([32.]) - if hasattr(ckpt, 'names') and isinstance(ckpt.names, (list, tuple)): - ckpt.names = dict(enumerate(ckpt.names)) # convert to dict + ckpt.args = {k: v for k, v in args.items() if k in default_keys} + # Append model.append(ckpt.fuse().eval() if fuse and hasattr(ckpt, 'fuse') else ckpt.eval()) # model in eval mode # Module compatibility updates @@ -310,7 +312,7 @@ def attempt_load_weights(weights, device=None, inplace=True, fuse=True): if len(model) == 1: return model[-1] - # Return detection ensemble + # Return ensemble print(f'Ensemble created with {weights}\n') for k in 'names', 'nc', 'yaml': setattr(model, k, getattr(model[0], k)) diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py index f422ce3..857a46c 100644 --- a/ultralytics/yolo/engine/exporter.py +++ b/ultralytics/yolo/engine/exporter.py @@ -164,8 +164,8 @@ class Exporter: assert not self.args.dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic' # Checks - if self.args.batch_size == 16: - self.args.batch_size = 1 # TODO: resolve batch_size 16 default in config.yaml + # if self.args.batch_size == model.args['batch_size']: # user has not modified training batch_size + self.args.batch_size = 1 self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2) # check image size if self.args.optimize: assert self.device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu' @@ -778,7 +778,7 @@ def export(cfg): if Path(cfg.model).suffix == '.yaml': model = DetectionModel(cfg.model) elif Path(cfg.model).suffix == '.pt': - model = attempt_load_weights(cfg.model) + model = attempt_load_weights(cfg.model, fuse=True) else: TypeError(f'Unsupported model type {cfg.model}') exporter(model=model) diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 0712f42..728eda7 100644 --- a/ultralytics/yolo/engine/model.py +++ b/ultralytics/yolo/engine/model.py @@ -77,13 +77,12 @@ class YOLO: Args: weights (str): model checkpoint to be loaded """ - self.ckpt = torch.load(weights, map_location="cpu") - self.task = self.ckpt["train_args"]["task"] - self.overrides = dict(self.ckpt["train_args"]) + self.model = attempt_load_weights(weights) + self.task = self.model.args["task"] + self.overrides = self.model.args self.overrides["device"] = '' # reset device self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = \ self._guess_ops_from_task(self.task) - self.model = attempt_load_weights(weights, fuse=False) def reset(self): """ @@ -189,7 +188,7 @@ class YOLO: raise AttributeError("dataset not provided! Please define `data` in config.yaml or pass as an argument.") self.trainer = self.TrainerClass(overrides=overrides) - self.trainer.model = self.trainer.load_model(weights=self.ckpt, + self.trainer.model = self.trainer.load_model(weights=self.model, model_cfg=self.model.yaml if self.task != "classify" else None) self.model = self.trainer.model # override here to save memory diff --git a/ultralytics/yolo/engine/validator.py b/ultralytics/yolo/engine/validator.py index d4f36b5..61cf6cf 100644 --- a/ultralytics/yolo/engine/validator.py +++ b/ultralytics/yolo/engine/validator.py @@ -106,6 +106,9 @@ class BaseValidator: data = check_dataset_yaml(self.args.data) else: data = check_dataset(self.args.data) + + if self.device.type == 'cpu': + self.args.workers = 0 # faster CPU val as time dominated by inference, not dataloading self.dataloader = self.dataloader or \ self.get_dataloader(data.get("val") or data.set("test"), self.args.batch_size) diff --git a/ultralytics/yolo/utils/__init__.py b/ultralytics/yolo/utils/__init__.py index c788987..15541a5 100644 --- a/ultralytics/yolo/utils/__init__.py +++ b/ultralytics/yolo/utils/__init__.py @@ -271,19 +271,20 @@ def yaml_save(file='data.yaml', data=None): yaml.safe_dump({k: str(v) if isinstance(v, Path) else v for k, v in data.items()}, f, sort_keys=False) -def yaml_load(file='data.yaml'): +def yaml_load(file='data.yaml', append_filename=True): """ Load YAML data from a file. Args: file (str, optional): File name. Default is 'data.yaml'. + append_filename (bool): Add the YAML filename to the YAML dictionary. Default is True. Returns: dict: YAML data and file name. """ with open(file, errors='ignore') as f: # Add YAML filename to dict and return - return {**yaml.safe_load(f), 'yaml_file': str(file)} + return {**yaml.safe_load(f), 'yaml_file': str(file)} if append_filename else yaml.safe_load(f) def get_settings(file=USER_CONFIG_DIR / 'settings.yaml'): diff --git a/ultralytics/yolo/v8/detect/train.py b/ultralytics/yolo/v8/detect/train.py index efc1296..db760d5 100644 --- a/ultralytics/yolo/v8/detect/train.py +++ b/ultralytics/yolo/v8/detect/train.py @@ -54,7 +54,7 @@ class DetectionTrainer(BaseTrainer): self.model.names = self.data["names"] def load_model(self, model_cfg=None, weights=None, verbose=True): - model = DetectionModel(model_cfg or weights["model"].yaml, ch=3, nc=self.data["nc"], verbose=verbose) + model = DetectionModel(model_cfg or weights.yaml, ch=3, nc=self.data["nc"], verbose=verbose) if weights: model.load(weights, verbose) return model diff --git a/ultralytics/yolo/v8/segment/train.py b/ultralytics/yolo/v8/segment/train.py index 37284d3..f8f146d 100644 --- a/ultralytics/yolo/v8/segment/train.py +++ b/ultralytics/yolo/v8/segment/train.py @@ -17,7 +17,7 @@ from ultralytics.yolo.utils.torch_utils import de_parallel class SegmentationTrainer(v8.detect.DetectionTrainer): def load_model(self, model_cfg=None, weights=None, verbose=True): - model = SegmentationModel(model_cfg or weights["model"].yaml, ch=3, nc=self.data["nc"], verbose=verbose) + model = SegmentationModel(model_cfg or weights.yaml, ch=3, nc=self.data["nc"], verbose=verbose) if weights: model.load(weights, verbose) return model