From 7bccd72d797020f3c39ba0b98cbb242466361b26 Mon Sep 17 00:00:00 2001 From: Christoph Gerum Date: Mon, 10 May 2021 17:18:58 +0200 Subject: [PATCH 1/2] Allow to manual selection of device for torchhub models --- hubconf.py | 45 +++++++++++++++++++++++++-------------------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/hubconf.py b/hubconf.py index 898b63b5d0b1..dabbfc4a2065 100644 --- a/hubconf.py +++ b/hubconf.py @@ -18,7 +18,7 @@ check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop')) -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: @@ -28,6 +28,7 @@ def create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbos 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 @@ -51,7 +52,11 @@ def create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbos 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 + if device is None: + device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available + elif isinstance(device, str): + device = torch.device(device) + return model.to(device) except Exception as e: @@ -60,49 +65,49 @@ def create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbos 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__': From 767c502719838f6cadf0ad06fa5e70de40a596e5 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 16 May 2021 16:24:28 +0200 Subject: [PATCH 2/2] single line device nested torch.device(torch.device(device)) ok --- hubconf.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/hubconf.py b/hubconf.py index 8ad8c9e263c3..f74e70c85a65 100644 --- a/hubconf.py +++ b/hubconf.py @@ -51,11 +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 - if device is None: - device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available - elif isinstance(device, str): - device = torch.device(device) - + 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: