Pytorch ModuleNotFoundError: No module named ‘models‘

问题

今天，在使用pytorch深度学习框架训练出来的模型文件，在另外的工程中使用，碰到如下的错误：

(shennong) ubuntu@ubuntu:~/zy/medicinebox_detector$ python interface_demo.py
Using torch 1.7.1 CUDA:0 (GeForce RTX 2080 Ti, 11016MB)

Traceback (most recent call last):
File “interface_demo.py”, line 13, in
detector = MedicineBoxDetector(save_img=False, weights=’/home/ubuntu/ModelWeight/yolov5x1280fillPolyBaseManual.pt’)
File “/home/ubuntu/zy/medicinebox_detector/medicine_detection/medicinebox_detector.py”, line 35, in init
self.model = attempt_load(weights, map_location=self.device) # load FP32 model
File “/home/ubuntu/zy/medicinebox_detector/medicine_detection/models/experimental.py”, line 118, in attempt_load
model.append(torch.load(w, map_location=map_location)[‘model’].float().fuse().eval()) # load FP32 model
File “/home/ubuntu/.conda/envs/shennong/lib/python3.6/site-packages/torch/serialization.py”, line 594, in load
return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
File “/home/ubuntu/.conda/envs/shennong/lib/python3.6/site-packages/torch/serialization.py”, line 853, in _load
result = unpickler.load()
ModuleNotFoundError: No module named ‘models’

下面是我的models所在的目录

├── interface_demo.py
…
├── medicine_detection
…
│ ├── models
│ │ ├── common.py
│ │ ├── experimental.py
│ │ ├── export.py
…

解决方法

这个问题在https://github.com/ultralytics/yolov5/issues/353这个地方已经解决
我参考

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

我把interface_demo.py这个文件加如下行，问题解决：

import sys
sys.path.insert(0, '/home/ubuntu/zy/medicinebox_detector/medicine_detection')