Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.

hubconf.py 3.0KB

il y a 4 ans
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. """File for accessing YOLOv5 via PyTorch Hub https://pytorch.org/hub/
  2. Usage:
  3. import torch
  4. model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, channels=3, classes=80)
  5. """
  6. dependencies = ['torch', 'yaml']
  7. import os
  8. import torch
  9. from models.yolo import Model
  10. from utils import google_utils
  11. def create(name, pretrained, channels, classes):
  12. """Creates a specified YOLOv5 model
  13. Arguments:
  14. name (str): name of model, i.e. 'yolov5s'
  15. pretrained (bool): load pretrained weights into the model
  16. channels (int): number of input channels
  17. classes (int): number of model classes
  18. Returns:
  19. pytorch model
  20. """
  21. config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) # model.yaml path
  22. model = Model(config, channels, classes)
  23. if pretrained:
  24. ckpt = '%s.pt' % name # checkpoint filename
  25. google_utils.attempt_download(ckpt) # download if not found locally
  26. state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32
  27. state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
  28. model.load_state_dict(state_dict, strict=False) # load
  29. return model
  30. def yolov5s(pretrained=False, channels=3, classes=80):
  31. """YOLOv5-small model from https://github.com/ultralytics/yolov5
  32. Arguments:
  33. pretrained (bool): load pretrained weights into the model, default=False
  34. channels (int): number of input channels, default=3
  35. classes (int): number of model classes, default=80
  36. Returns:
  37. pytorch model
  38. """
  39. return create('yolov5s', pretrained, channels, classes)
  40. def yolov5m(pretrained=False, channels=3, classes=80):
  41. """YOLOv5-medium model from https://github.com/ultralytics/yolov5
  42. Arguments:
  43. pretrained (bool): load pretrained weights into the model, default=False
  44. channels (int): number of input channels, default=3
  45. classes (int): number of model classes, default=80
  46. Returns:
  47. pytorch model
  48. """
  49. return create('yolov5m', pretrained, channels, classes)
  50. def yolov5l(pretrained=False, channels=3, classes=80):
  51. """YOLOv5-large model from https://github.com/ultralytics/yolov5
  52. Arguments:
  53. pretrained (bool): load pretrained weights into the model, default=False
  54. channels (int): number of input channels, default=3
  55. classes (int): number of model classes, default=80
  56. Returns:
  57. pytorch model
  58. """
  59. return create('yolov5l', pretrained, channels, classes)
  60. def yolov5x(pretrained=False, channels=3, classes=80):
  61. """YOLOv5-xlarge model from https://github.com/ultralytics/yolov5
  62. Arguments:
  63. pretrained (bool): load pretrained weights into the model, default=False
  64. channels (int): number of input channels, default=3
  65. classes (int): number of model classes, default=80
  66. Returns:
  67. pytorch model
  68. """
  69. return create('yolov5x', pretrained, channels, classes)