From 94a1e71c6d00b9067b8b11e5f4f53f2a1bb2ffb8 Mon Sep 17 00:00:00 2001 From: Christoph Gerum Date: Sun, 16 May 2021 17:41:26 +0200 Subject: [PATCH] Add `device` argument to PyTorch Hub models (#3104) * Allow to manual selection of device for torchhub models * single line device nested torch.device(torch.device(device)) ok Co-authored-by: Glenn Jocher --- hubconf.py | 41 +++++++++++++++++++++-------------------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/hubconf.py b/hubconf.py index 3b3dfe0e9e23..f74e70c85a65 100644 --- a/hubconf.py +++ b/hubconf.py @@ -8,7 +8,7 @@ import torch -def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): +def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): """Creates a specified YOLOv5 model Arguments: @@ -18,6 +18,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo classes (int): number of model classes autoshape (bool): apply YOLOv5 .autoshape() wrapper to model verbose (bool): print all information to screen + device (str, torch.device, None): device to use for model parameters Returns: YOLOv5 pytorch model @@ -50,7 +51,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo model.names = ckpt['model'].names # set class names attribute if autoshape: model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS - device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available + device = select_device('0' if torch.cuda.is_available() else 'cpu') if device is None else torch.device(device) return model.to(device) except Exception as e: @@ -59,49 +60,49 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo raise Exception(s) from e -def custom(path='path/to/model.pt', autoshape=True, verbose=True): +def custom(path='path/to/model.pt', autoshape=True, verbose=True, device=None): # YOLOv5 custom or local model - return _create(path, autoshape=autoshape, verbose=verbose) + return _create(path, autoshape=autoshape, verbose=verbose, device=device) -def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): +def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-small model https://github.com/ultralytics/yolov5 - return _create('yolov5s', pretrained, channels, classes, autoshape, verbose) + return _create('yolov5s', pretrained, channels, classes, autoshape, verbose, device) -def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): +def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-medium model https://github.com/ultralytics/yolov5 - return _create('yolov5m', pretrained, channels, classes, autoshape, verbose) + return _create('yolov5m', pretrained, channels, classes, autoshape, verbose, device) -def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): +def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-large model https://github.com/ultralytics/yolov5 - return _create('yolov5l', pretrained, channels, classes, autoshape, verbose) + return _create('yolov5l', pretrained, channels, classes, autoshape, verbose, device) -def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): +def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-xlarge model https://github.com/ultralytics/yolov5 - return _create('yolov5x', pretrained, channels, classes, autoshape, verbose) + return _create('yolov5x', pretrained, channels, classes, autoshape, verbose, device) -def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): +def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-small-P6 model https://github.com/ultralytics/yolov5 - return _create('yolov5s6', pretrained, channels, classes, autoshape, verbose) + return _create('yolov5s6', pretrained, channels, classes, autoshape, verbose, device) -def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): +def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5 - return _create('yolov5m6', pretrained, channels, classes, autoshape, verbose) + return _create('yolov5m6', pretrained, channels, classes, autoshape, verbose, device) -def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): +def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-large-P6 model https://github.com/ultralytics/yolov5 - return _create('yolov5l6', pretrained, channels, classes, autoshape, verbose) + return _create('yolov5l6', pretrained, channels, classes, autoshape, verbose, device) -def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): +def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5 - return _create('yolov5x6', pretrained, channels, classes, autoshape, verbose) + return _create('yolov5x6', pretrained, channels, classes, autoshape, verbose, device) if __name__ == '__main__':