mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Simplified PyTorch hub for custom models (#1677)

pull/1700/head
Glenn Jocher 2020-12-12 11:37:53 -08:00 committed by GitHub
parent 87ca35b922
commit f7a923b446
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
hubconf.py
View File

@ -106,19 +106,18 @@ def yolov5x(pretrained=False, channels=3, classes=80):
return create('yolov5x', pretrained, channels, classes)
def custom(model='path/to/model.pt'):
def custom(path_or_model='path/to/model.pt'):
"""YOLOv5-custom model from https://github.com/ultralytics/yolov5
Arguments (3 format options):
model (str): 'path/to/model.pt'
model (dict): torch.load('path/to/model.pt')
model (nn.Module): 'torch.load('path/to/model.pt')['model']
Arguments (3 options):
path_or_model (str): 'path/to/model.pt'
path_or_model (dict): torch.load('path/to/model.pt')
path_or_model (nn.Module): torch.load('path/to/model.pt')['model']
Returns:
pytorch model
"""
if isinstance(model, str):
model = torch.load(model) # load checkpoint
model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint
if isinstance(model, dict):
model = model['model'] # load model