pytorch加载模型torch.load()遇到问题ModuleNotFoundError: No module named models

问题描述：

在加载自己训练的模型文件时，报错ModuleNotFoundError: No module named ‘models’。

import torch
model_tmp = torch.load('./yolov5s.pt')
print(model_tmp)

问题分析：

写load的这个文件要放到具有yolov5工程的文件夹里运行，因为它会寻找原生yolo5工程下的models文件夹，里面还会有yolo.py等文件存在，这就是这个找不到”models”的由来。

问题解决：

将代码加载权重文件所在的目录，新建models文件，然后把模型.py放到该目录，就没报错了。

虽然torch.save()既保存了模型又保存了权重参数文件，但是还需要将模型放到torch加载的地方。