RuntimeError: Error(s) in loading state_dict for AutoShape

[TheTechPny]

I cloned the yolov5 repo and trained it using a custom dataset I created. The training went fine for the quality of my dataset however I can't seem to load the weights I've just created.

python train.py --img 416 --batch 4 --epochs 100 --data mytest.yaml --cfg yolov5s.yaml --name test_results --cache

` def init(self): self.lock = Lock()

load the trained model

    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=False, classes=2)
    checkpoint_ = torch.load('weights/best.pt')['model']
    model.load_state_dict(checkpoint_.state_dict())`

When attempting to load the weights with the above code I get this error https://pastebin.com/raw/KX5nkHiA

[github-actions[bot]]

👋 Hello @TheTechPny, thank you for your interest in YOLOv5 🚀! Please visit our ⭐️ Tutorials to get started, where you can find quickstart guides for simple tasks like Custom Data Training all the way to advanced concepts like Hyperparameter Evolution.

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 training ❓ Question, please provide as much information as possible, including dataset images, training logs, screenshots, and a public link to online W&B logging if available.

For business inquiries or professional support requests please visit https://ultralytics.com or email Glenn Jocher at glenn.jocher@ultralytics.com.

Requirements

Python>=3.6.0 with all requirements.txt installed including PyTorch>=1.7. To get started:

$ git clone https://github.com/ultralytics/yolov5
$ cd yolov5
$ pip install -r requirements.txt

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Google Colab and Kaggle notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training (train.py), validation (val.py), inference (detect.py) and export (export.py) on MacOS, Windows, and Ubuntu every 24 hours and on every commit.

[TheTechPny]

I also trained the model in Google Colab with the same dataset and I couldn't locally load the weights either

[glenn-jocher]

@TheTechPny see PyTorch Hub tutorial for loading custom YOLOv5 models.

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

[TheTechPny]

I already have my weights. I've trained the model properly, the problem I'm having is with loading the weights. I don't see anything in the tutorial about properly loading weights and I've implemented the pytorch load_state_dict code appropriately, I'm not sure what the issue is.

    # load the trained model
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=False)
    checkpoint_ = torch.load('weights/best.pt')['model']
    self.model = model.load_state_dict(checkpoint_.state_dict())
    self.model.conf = 0.2  # confidence threshold (0-1)

https://i.imgur.com/IK0mdGC.png

[glenn-jocher]

@TheTechPny your workflow is incorrect. See PyTorch Hub tutorial for correct workflow.

[TheTechPny]

@TheTechPny your workflow is incorrect. See PyTorch Hub tutorial for correct workflow.

Sorry for wasting your time!

For anyone that stumbles across this same problem, look at the 'Custom Models' section of the PyTorch Hub tutorial

model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/best.pt') # default model = torch.hub.load('path/to/yolov5', 'custom', path='path/to/best.pt', source='local') # local repo