Simplified PyTorch hub custom models

[glenn-jocher]

This PR simplifies the process of loading a custom YOLOv5-based model (of any architecture) in PyTorch Hub. The new syntax is very simple:

import torch

model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='path/to/custom_model.pt')
model = model.autoshape()  # for PIL/cv2/np inputs and NMS

# OR if a custom model is already loaded in the python workspace:
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model=custom_model)
model = model.autoshape()  # for PIL/cv2/np inputs and NMS

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

New support for loading custom YOLOv5 models directly from user-specified paths.

📊 Key Changes

• Added custom() function to hubconf.py allowing users to load their own trained YOLOv5 models.
• Supports three input formats:
  • A string path to the model .pt file.
  • A dictionary object obtained from torch.load().
  • An nn.Module object from the loaded checkpoint.
• Automatically transfers the model to the same device as its parameters.
• Loads the model's state_dict and class names (names).
• Example usage commented out in the __main__ block.

🎯 Purpose & Impact

• 🚀 Enables users to easily deploy their custom-trained YOLOv5 models for inference.
• 💻 Improves flexibility by supporting different methods for providing the model.
• 🔌 Potentially expands the use cases of YOLOv5 to custom datasets and specialized applications.
• 🌐 Non-expert users benefit from a straightforward way to implement their own models, improving the overall user experience.

[NanoCode012]

Hello, is there a reason why auto shape is not done automatically(inside the functions)? Is there a case where users don’t need auto shape?

[glenn-jocher]

@NanoCode012 that's a good question. The only reason not to use .autoshape() would be if you want to use the model for training (with your own custom training code) rather than for inference. I'd assume very few people are doing this though, as most would probably just use this repository to train and then place the model in their project for inference.

The problem is that the autoshape action is not reversible, mostly because it contains a call to .fuse(), so if the model is provided by default with autoshape then it would not be possible to train with it.

[NanoCode012]

@glenn-jocher , perhaps you can add an argument like autoshape=False or trainable=False and set autoshape to True by default?

[glenn-jocher]

@NanoCode012 actually that's a good idea. Then instead of 99% of the people being confused it will invert and the 1% that want to train will be confused, which should work a lot better for the majority. Ok I'll add a TODO here for that.

TODO: think about enabling autoshape by default on all PyTorch Hub model.

[glenn-jocher]

Removing TODO as PR #1692 is now open.