How can I move and run the classification model out of YOLOV5 codebase?

[AI-Passionner]

Search before asking

• I have searched the YOLOv5 issues and discussions and found no similar questions.

Question

It seems that loading the classification model can't leave out of the YOLOV5-Classifier code base. When loading the model,
model = torch.load('path/to/best.pt', map_location=torch.device('cpu'))['model'].float(), it throws error like this
"""
module = self._system_import(name, *args, **kwargs)
ModuleNotFoundError: No module named 'models'
"""
Once the folders of models and utils are copied to a new project directory, the loading works.

I am asking if there is a way I can load the classification model and use it as a stand-alone and for example, converting to TensorFlow like the object detection codebase.
As I know, the YOLOV5 object detection doesn't have this problem.

Thanks.

Additional

No response

[glenn-jocher]

@AI-Passionner this is a good question. In general for any pytorch model you'll need to have the modules available in the workspace prior to loading. PyTorch Hub accomplishes this in the background by downloading and caching the repo in a separate directory.

To point PyTorch Hub to a specific branch in a repo, i.e. the classifier branch in the YOLOv5 repo you can try this, though the classifier should be merged into master as part of the upcoming v6.2 release in the next few weeks.

model = torch.hub.load('ultralytics/yolov5:classifier', 'custom', path='path/to/model.pt')  # local model

[AI-Passionner]

@glenn-jocher Thank you. Let me try. I bet this is what I am looking for. Really appreciate your help.
Tiger

[AI-Passionner]

The classification model was trained as yolov5s. I copied the model file to the new project and. But it didn't work when loading the new model as your suggested scripts.
model = torch.hub.load('ultralytics/yolov5:classifier', 'custom', path='weights/yolov5s_exp30/weights/best.pt')

I threw the error message like this. Adding force_reload=True didn't work. Deleting the Cache didn't work. Any suggestion?
Do I need to specify the custom in the script?

""""
Fusing layers...
YOLOv5s summary: 117 layers, 4169250 parameters, 0 gradients, 10.4 GFLOPs
Adding AutoShape...
Traceback (most recent call last):
File "C:\Users\tiger.cao/.cache\torch\hub\ultralytics_yolov5_classifier\hubconf.py", line 62, in _create
return model.to(device)
File "C:\Users\tiger.cao\anaconda3\envs\Yolov5\lib\site-packages\torch\nn\modules\module.py", line 927, in to
return self._apply(convert)
File "C:\Users\tiger.cao/.cache\torch\hub\ultralytics_yolov5_classifier\models\common.py", line 571, in _apply
m.stride = fn(m.stride)
File "C:\Users\tiger.cao\anaconda3\envs\Yolov5\lib\site-packages\torch\nn\modules\module.py", line 1207, in getattr
raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'Classify' object has no attribute 'stride'
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "C:\Users\tiger.cao\anaconda3\envs\Yolov5\lib\site-packages\IPython\core\interactiveshell.py", line 3398, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "", line 1, in <cell line: 1>
model = torch.hub.load('ultralytics/yolov5:classifier', 'custom', force_reload=True, path='weights/yolov5s_exp30/weights/best.pt') # local model
File "C:\Users\tiger.cao\anaconda3\envs\Yolov5\lib\site-packages\torch\hub.py", line 540, in load
model = _load_local(repo_or_dir, model, *args, **kwargs)
File "C:\Users\tiger.cao\anaconda3\envs\Yolov5\lib\site-packages\torch\hub.py", line 569, in _load_local
model = entry(*args, **kwargs)
File "C:\Users\tiger.cao/.cache\torch\hub\ultralytics_yolov5_classifier\hubconf.py", line 72, in custom
return _create(path, autoshape=autoshape, verbose=_verbose, device=device)
File "C:\Users\tiger.cao/.cache\torch\hub\ultralytics_yolov5_classifier\hubconf.py", line 67, in _create
raise Exception(s) from e
Exception: 'Classify' object has no attribute 'stride'. Cache may be out of date, try force_reload=True or see https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading for help.
""""

[glenn-jocher]

@AI-Passionner got it. It looks like we need to add some classifier-specific code to hubconf.py to handle these models differently, which makes sense. I'll add a TODO for this.

[AI-Passionner]

@glenn-jocher Thank you.
Currently, I copied the folders models/ and utils/ with a little bit of cleaning to the new project. It works.

[glenn-jocher]

@AI-Passionner I think this should work now for newly trained models of type models.yolo.ClassificationModel, (but not for older models of type models.yolo.Model).

I'll run some tests to verify today.

[AI-Passionner]

@glenn-jocher Thank you. I will run the train again and have a try.
Tiger

[glenn-jocher]

@AI-Passionner I've confirmed PyTorch Hub loading works in a test just now on Google Colab. I uploaded a recent ImageNet trained YOLOv5m-cls model (about 75 top-1 accuracy after 90 epochs) and then requested it like this. This does not need a local git clone of YOLOv5.
https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5m-cls.pt

import torch

model = torch.hub.load('ultralytics/yolov5:classifier', 'custom', 'yolov5m-cls.pt')  # or use local cls model
print(type(model))

EDIT: Once we release 6.2 with built-in classification support then you no longer need to point torch.hub.load() to the classifier branch, i.e. you can use torch.hub.load('ultralytics/yolov5', ...)