Browse Source

Scope imports for torch.hub.list() improvement (#3144)

modifyDataloader
Glenn Jocher GitHub 3 years ago
parent
commit
1935266951
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 15 additions and 14 deletions
  1. +15
    -14
      hubconf.py

+ 15
- 14
hubconf.py View File



import torch import torch


from models.yolo import Model, attempt_load
from utils.general import check_requirements, set_logging from utils.general import check_requirements, set_logging
from utils.google_utils import attempt_download
from utils.torch_utils import select_device


dependencies = ['torch', 'yaml'] dependencies = ['torch', 'yaml']
check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop')) 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):
"""Creates a specified YOLOv5 model """Creates a specified YOLOv5 model


Arguments: Arguments:
Returns: Returns:
YOLOv5 pytorch model YOLOv5 pytorch model
""" """
from models.yolo import Model, attempt_load
from utils.google_utils import attempt_download
from utils.torch_utils import select_device

set_logging(verbose=verbose) set_logging(verbose=verbose)
fname = Path(name).with_suffix('.pt') # checkpoint filename fname = Path(name).with_suffix('.pt') # checkpoint filename
try: try:


def custom(path='path/to/model.pt', autoshape=True, verbose=True): def custom(path='path/to/model.pt', autoshape=True, verbose=True):
# YOLOv5 custom or local model # YOLOv5 custom or local model
return create(path, autoshape=autoshape, verbose=verbose)
return _create(path, autoshape=autoshape, verbose=verbose)




def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-small model https://github.com/ultralytics/yolov5 # YOLOv5-small model https://github.com/ultralytics/yolov5
return create('yolov5s', pretrained, channels, classes, autoshape, verbose)
return _create('yolov5s', pretrained, channels, classes, autoshape, verbose)




def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-medium model https://github.com/ultralytics/yolov5 # YOLOv5-medium model https://github.com/ultralytics/yolov5
return create('yolov5m', pretrained, channels, classes, autoshape, verbose)
return _create('yolov5m', pretrained, channels, classes, autoshape, verbose)




def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-large model https://github.com/ultralytics/yolov5 # YOLOv5-large model https://github.com/ultralytics/yolov5
return create('yolov5l', pretrained, channels, classes, autoshape, verbose)
return _create('yolov5l', pretrained, channels, classes, autoshape, verbose)




def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-xlarge model https://github.com/ultralytics/yolov5 # YOLOv5-xlarge model https://github.com/ultralytics/yolov5
return create('yolov5x', pretrained, channels, classes, autoshape, verbose)
return _create('yolov5x', pretrained, channels, classes, autoshape, verbose)




def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-small-P6 model https://github.com/ultralytics/yolov5 # 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)




def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5 # 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)




def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-large-P6 model https://github.com/ultralytics/yolov5 # 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)




def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5 # 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)




if __name__ == '__main__': if __name__ == '__main__':
model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained
model = _create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained
# model = custom(path='path/to/model.pt') # custom # model = custom(path='path/to/model.pt') # custom


# Verify inference # Verify inference

Loading…
Cancel
Save