Glenn Jocher
GitHub
3 years ago
1 changed files with 6 additions and 7 deletions
+ 6
- 7
hubconf.py
View File
| 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): | |||||
| 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: | Returns: | ||||
| pytorch model | 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): | if isinstance(model, dict): | ||||
| model = model['model'] # load model | model = model['model'] # load model | ||||
Loading…