Update train.py to save and load only the model parameters

[mmenbawy]

🚀 Feature

Update train.py to save only the state dicts and not whole model objects.

Motivation

It's extremely frustrating to retrain the model to produce a new set of weights and then find up that the weights can not be used in a different directory structure.
#22
and per pytorch/pytorch#3678 torch recommends saving only the state dicts and not whole model objects to avoid having serialized data bound to the specific classes and the exact directory structure used.

Pitch

Change

yolov5/train.py

Line 353 in c67e722

torch.save(ckpt, last)

to torch.save(ckpt.state_dict(), last)

Alternatives

Additional context

[github-actions]

Hello @Mostafa-Elmenbawy, thank you for your interest in our work! Please visit our Custom Training Tutorial to get started, and see our Jupyter Notebook , Docker Image, and Google Cloud Quickstart Guide for example environments.

If this is a bug report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom model or data training question, please note Ultralytics does not provide free personal support. As a leader in vision ML and AI, we do offer professional consulting, from simple expert advice up to delivery of fully customized, end-to-end production solutions for our clients, such as:

• Cloud-based AI systems operating on hundreds of HD video streams in realtime.
• Edge AI integrated into custom iOS and Android apps for realtime 30 FPS video inference.
• Custom data training, hyperparameter evolution, and model exportation to any destination.

For more information please visit https://www.ultralytics.com.

[glenn-jocher]

@Mostafa-Elmenbawy yes, I understand, thanks for the feedback.

We saved separate weights and cfg files in our older https://github.com/ultralytics/yolov3 repo, but we saw that this generated confusion among users, who would inadvertently attempt to pair incompatible cfg and weights files, and of course also added complications when exporting, since two arguments were required rather than one whenever a model needed to be specified.

So with YOLOv5 we created simplified standalone models that are self contained with all info required (config, weights, anchors, class names, etc.). Inference with yolov3 for example requires you to pass the --cfg, the --weights, and the --names. YOLOv5 inference requires only --weights, an argument reduction of 3x.

The easiest way to use YOLOv5 independent of this repo is PyTorch Hub. See https://docs.ultralytics.com/yolov5

[glenn-jocher]

@Mostafa-Elmenbawy see #1153 (comment) for our awesome new standalone YOLOv5 functionality.

You pass in a PIL image, you get out detections, everything handled automatically, and no dependency on cloning the ultralytics/yolov5 repo.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[Sagor-Saha]

@Mostafa-Elmenbawy you saved the state_dict of the model while training. How to load the model during prediction of images? Code?