Cannot load model using torch.hub

[matthewfranglen]

🐛 Bug

I cannot use your hubconf.py to load the model.

To Reproduce (REQUIRED)

Input:

import torch.hub

# all of these fail
model = torch.hub.load("ultralytics/yolov5", "yolov5s", pretrained=True)
model = torch.hub.load("ultralytics/yolov5", "yolov5m", pretrained=True)
model = torch.hub.load("ultralytics/yolov5", "yolov5l", pretrained=True)

Output:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-6-f8aae7484cc6> in <module>
----> 1 model = torch.hub.load("ultralytics/yolov5", "yolov5l", pretrained=True)

~/.cache/pypoetry/virtualenvs/dl-ultralytics-yolo-kIMoBkIY-py3.7/lib/python3.7/site-packages/torch/hub.py in load(github, model, *args, **kwargs)
    367     entry = _load_entry_from_hubconf(hub_module, model)
    368 
--> 369     model = entry(*args, **kwargs)
    370 
    371     sys.path.remove(repo_dir)

~/.cache/torch/hub/ultralytics_yolov5_master/hubconf.py in yolov5l(pretrained, channels, classes)
     78         pytorch model
     79     """
---> 80     return create('yolov5l', pretrained, channels, classes)
     81 
     82 

~/.cache/torch/hub/ultralytics_yolov5_master/hubconf.py in create(name, pretrained, channels, classes)
     34         google_utils.attempt_download(ckpt)  # download if not found locally
     35         state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict()  # to FP32
---> 36         state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape}  # filter
     37         model.load_state_dict(state_dict, strict=False)  # load
     38     return model

~/.cache/torch/hub/ultralytics_yolov5_master/hubconf.py in <dictcomp>(.0)
     34         google_utils.attempt_download(ckpt)  # download if not found locally
     35         state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict()  # to FP32
---> 36         state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape}  # filter
     37         model.load_state_dict(state_dict, strict=False)  # load
     38     return model

KeyError: 'model.2.cv1.conv.weight'

Expected behavior

The model should be created and the pretrained weights should be loaded.

[github-actions[bot]]

Hello @matthewfranglen, 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 that 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]

@matthewfranglen this issue is not reproducible. Open a blank colab notebook, update pyyaml, restart and import the model. Everything works correctly here. If you are seeing errors you may want to delete the local cache that is being used and retry.

[matthewfranglen]

Deleting the cache file fixed the problem. Thank you for the speedy response.

[glenn-jocher]

@matthewfranglen yes, this is happening because the repo is changing fast, and we are introducing breaking changes in the process.

I wonder if we could wrap the entire hub loading function in a try except to notify future users to delete their cache when a failure occurs.

[matthewfranglen]

I think that would help a lot

[glenn-jocher]

All done. Relevant code wrapped in try except clause now. Seems like a great idea. https://github.com/ultralytics/yolov5/blob/a5818b280edcc983cb6bec626f36db50a5221f68/hubconf.py#L31-L43

[hysaint]

Dear esteemed author: yeah i delete the cache ,however the commad continuning leave the error message Exception: Cache maybe be out of date, try force_reload=True. See https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading for help. I try to add the force_reload=True argument but still not work.Could u plz help me to find out the solution？ Thanks

[glenn-jocher]

@hysaint YOLOv5 hub works correctly. Paste the following code into a new Colab notebook:

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# Images
dir = 'https://github.com/ultralytics/yolov5/raw/master/data/images/'
imgs = [dir + f for f in ('zidane.jpg', 'bus.jpg')]  # batched list of images

# Inference
results = model(imgs)
results.print()  # or .show(), .save()

You'll see this:

[hysaint]

@glenn-jocher very thankful .it works great