From e486a2e0ccccb7c6115456b5adc547d1576066f5 Mon Sep 17 00:00:00 2001 From: johnohagan <86861886+johnohagan@users.noreply.github.com> Date: Tue, 6 Jul 2021 21:39:57 +1000 Subject: [PATCH 1/2] Create hubconf.py --- hubconf.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/hubconf.py b/hubconf.py index 2de71d617f1e..cc4a9bc70292 100644 --- a/hubconf.py +++ b/hubconf.py @@ -33,7 +33,8 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo check_requirements(requirements=Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python')) set_logging(verbose=verbose) - + + name = Path(__file__).parent.parent / 'checkpoints' / name fname = Path(name).with_suffix('.pt') # checkpoint filename try: device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device) From 3ed7852aaf5ee0dfb0391644d4952af4b1440c87 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 7 Jul 2021 13:17:55 +0200 Subject: [PATCH 2/2] Add save_dir variable --- hubconf.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/hubconf.py b/hubconf.py index cc4a9bc70292..df268b18d177 100644 --- a/hubconf.py +++ b/hubconf.py @@ -4,9 +4,12 @@ import torch model = torch.hub.load('ultralytics/yolov5', 'yolov5s') """ +from pathlib import Path import torch +FILE = Path(__file__).absolute() + def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): """Creates a specified YOLOv5 model @@ -23,29 +26,26 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo Returns: YOLOv5 pytorch model """ - from pathlib import Path - from models.yolo import Model, attempt_load from utils.general import check_requirements, set_logging from utils.google_utils import attempt_download from utils.torch_utils import select_device - check_requirements(requirements=Path(__file__).parent / 'requirements.txt', - exclude=('tensorboard', 'thop', 'opencv-python')) + check_requirements(requirements=FILE.parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python')) set_logging(verbose=verbose) - - name = Path(__file__).parent.parent / 'checkpoints' / name - fname = Path(name).with_suffix('.pt') # checkpoint filename + + save_dir = Path('') if str(name).endswith('.pt') else FILE.parent + path = (save_dir / name).with_suffix('.pt') # checkpoint path try: device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device) if pretrained and channels == 3 and classes == 80: - model = attempt_load(fname, map_location=device) # download/load FP32 model + model = attempt_load(path, map_location=device) # download/load FP32 model else: cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path model = Model(cfg, channels, classes) # create model if pretrained: - ckpt = torch.load(attempt_download(fname), map_location=device) # load + ckpt = torch.load(attempt_download(path), map_location=device) # load msd = model.state_dict() # model state_dict csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape} # filter