参考博客:
ModuleNotFoundError: No module named ‘models‘解决torch.load问题【天坑】

保存与加载

使用 torch.save(model, “my_model.pth”) 命令可以保存整个模型。
这个保存/加载过程使用最直观的语法,涉及的代码最少。
以这种方式保存模型将使用Python的 pickle 模块保存整个model。
但是,在进行torch.load(“my_model.pth”)时,加载目录与保存目录要相同,这里的目录不是”my_model.pth”,而是项目中定义模型所涉及的目录。举个例子:

load-test工程中有model_1文件夹,model_1文件夹中有yolo.py模块。

yolo.py:

import torchclass MyNet(torch.nn.Module):def __init__(self):super(MyNet, self).__init__()# 第一句话,调用父类的构造函数self.conv1 = torch.nn.Conv2d(3, 32, 3, 1, 1)self.relu1 = torch.nn.ReLU()self.max_pooling1 = torch.nn.MaxPool2d(2, 1)self.conv2 = torch.nn.Conv2d(32, 32, 3, 1, 1)self.relu2 = torch.nn.ReLU()self.max_pooling2 = torch.nn.MaxPool2d(2, 1)self.dense1 = torch.nn.Linear(32 * 3 * 3, 128)self.dense2 = torch.nn.Linear(128, 10)def forward(self, x):x = self.conv1(x)x = self.relu1(x)x = self.max_pooling1(x)x = self.conv2(x)x = self.relu2(x)x = self.max_pooling2(x)x = self.dense1(x)x = self.dense2(x)return x

在train.py中导入模型并保存整个模型:

from model_1.yolo import MyNetimport torchnet = MyNet()torch.save(net, './weights/net.pt')# 保存

将net.pt保存在weights文件夹中。
然后在train.py中加载模型:

import torchmodel = torch.load('./weights/net.pt')

是成功的(也不需要 from model_1.yolo import MyNet )。

如果将model_1文件夹改为model_2文件夹,则出错。

出错的原因是导入模型时的目录与保存时的目录不一致。

解决方案

有时候需要将训练好的模型导入另一个工程中,但是该工程的文件夹与原文件夹相同且不方便更改,这时可以采取: 先加载模型,然后保存模型参数,再加载模型参数 的方法来正确加载模型。

先保持model_1文件夹不变。
train.py:

import torchmodel = torch.load('./weights/net.pt')torch.save(model.state_dict(), './weights/net_state_dict.pt')# 保存

将model_1文件夹改为model_2文件夹。
train.py:

from model_2.yolo import MyNetimport torchmodel = MyNet()state_dict = torch.load('./weights/net_state_dict.pt')model.load_state_dict(state_dict)

加载成功。