Train a custom dataset with PyThorch Hub

[alevangel]

I know that for training a custom dataset with train.py there is an extensive documentation, but how to dao that with PyTorch?

I only see from the doc on PyTorch that I have have to use this line:
model = torch.hub.load('ultralytics/yolov5', 'yolov5m', autoshape=False)

But then how to train the model?

[glenn-jocher]

@alevangel the command above will produce a pretrained YOLOv5m model, you provide the training code to train it.

See Train Custom Data tutorial below for a fully managed training solution.

YOLOv5 Tutorials

• Train Custom Data  🚀 RECOMMENDED
• Tips for Best Training Results  ☘️ RECOMMENDED
• Weights & Biases Logging  🌟 NEW
• Supervisely Ecosystem  🌟 NEW
• Multi-GPU Training
• PyTorch Hub  ⭐ NEW
• TorchScript, ONNX, CoreML Export 🚀
• Test-Time Augmentation (TTA)
• Model Ensembling
• Model Pruning/Sparsity
• Hyperparameter Evolution
• Transfer Learning with Frozen Layers  ⭐ NEW
• TensorRT Deployment

[alevangel]

But I only see train a custom dataset with a notebook using the train.py,
is possible to do that directly in PyCharm for istance?

Sorry if this is a dumb question

[glenn-jocher]

@alevangel PyCharm is a python IDE, naturally it's capable of everything a notebook can do.

[GTimothee]

@alevangel in this repo yolov5 gives you a train.py script that embed all the training functions etc.

When you call train.py the script loads the model, loads the dataset and execute a training function.

If you want to do it yourself, you can load the model yourself using pytorch hub as you mentioned, and then you have to load your dataset yourself, write the training function etc.

For more information and if you are a complete beginner, I would suggest following this tutorial. It explains how to fine tune a pytorch model. Then , just use the model defined with model = torch.hub.load('ultralytics/yolov5', 'yolov5m', autoshape=False) in the training loop of the tutorial.

Hope it helps.

[alevangel]

@GTimothee thanks a lot for this answer! Is what I was looking for, it helped a lot!