ModuleNotFoundError: No module named 'models'

hi yolov5, i met this problem when i try to use the model in my project. the question has solved but i think it's enough classical to open a new issue to describe it

repreduct

in yolov5 repo, the infer file is detect.py and the model is ./weights/yolov5s.pt. The complete detection code is as follows

python detect.py --source ./inference/images/ --weights yolov5s.pt --conf 0.4

i retrained the model and put it in my repo: ./weights/best.pt. and i put yolov5's models and utils in my repo

└── yolov5
    ├── detect.py
    ├── models
    ├── __pycache__
    └── utils

i use ./demo/test.py to predict and there is something wrong with the program

reason

refer to

torch.load() requires model module in the same folder #3678

ModuleNotFoundError: No module named 'models' #18325

Pytorch.load() error:No module named ‘model’

the key of this problem is Pickle need the same file structure between the calling file and the source file if use the following code to save model

torch.save(model, PATH)

i checked the source file train.py

            # Save last, best and delete
            torch.save(ckpt, last)
            if (best_fitness == fi) and not final_epoch:
                torch.save(ckpt, best)
            del ckpt

Yes, that's it

solution

use sys.path to include the source file. do it like this:

import sys
sys.path.insert(0, './yolov5')