Unverified 提交 61047a2b authored 作者: johnohagan's avatar johnohagan 提交者: GitHub

Save PyTorch Hub models to `/root/hub/cache/dir` (#3904)

* Create hubconf.py * Add save_dir variable Co-authored-by: 's avatarGlenn Jocher <glenn.jocher@ultralytics.com>
上级 33202b7f
...@@ -4,9 +4,12 @@ Usage: ...@@ -4,9 +4,12 @@ Usage:
import torch import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
""" """
from pathlib import Path
import torch import torch
FILE = Path(__file__).absolute()
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
"""Creates a specified YOLOv5 model """Creates a specified YOLOv5 model
...@@ -23,28 +26,26 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo ...@@ -23,28 +26,26 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
Returns: Returns:
YOLOv5 pytorch model YOLOv5 pytorch model
""" """
from pathlib import Path
from models.yolo import Model, attempt_load from models.yolo import Model, attempt_load
from utils.general import check_requirements, set_logging from utils.general import check_requirements, set_logging
from utils.google_utils import attempt_download from utils.google_utils import attempt_download
from utils.torch_utils import select_device from utils.torch_utils import select_device
check_requirements(requirements=Path(__file__).parent / 'requirements.txt', check_requirements(requirements=FILE.parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python'))
exclude=('tensorboard', 'thop', 'opencv-python'))
set_logging(verbose=verbose) set_logging(verbose=verbose)
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: try:
device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device) 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: 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: else:
cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path
model = Model(cfg, channels, classes) # create model model = Model(cfg, channels, classes) # create model
if pretrained: 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 msd = model.state_dict() # model state_dict
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 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 csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape} # filter
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论