Unverified 提交 f7a923b4 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

Simplified PyTorch hub for custom models (#1677)

上级 87ca35b9
...@@ -106,19 +106,18 @@ def yolov5x(pretrained=False, channels=3, classes=80): ...@@ -106,19 +106,18 @@ def yolov5x(pretrained=False, channels=3, classes=80):
return create('yolov5x', pretrained, channels, classes) 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 """YOLOv5-custom model from https://github.com/ultralytics/yolov5
Arguments (3 format options): Arguments (3 options):
model (str): 'path/to/model.pt' path_or_model (str): 'path/to/model.pt'
model (dict): torch.load('path/to/model.pt') path_or_model (dict): torch.load('path/to/model.pt')
model (nn.Module): 'torch.load('path/to/model.pt')['model'] path_or_model (nn.Module): torch.load('path/to/model.pt')['model']
Returns: Returns:
pytorch model pytorch model
""" """
if isinstance(model, str): model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint
model = torch.load(model) # load checkpoint
if isinstance(model, dict): if isinstance(model, dict):
model = model['model'] # load model model = model['model'] # load model
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论