在 PyTorch 中加载迁移学习模型进行推理的正确方法是啥？

Posted 2023-03-12

【中文标题】在 PyTorch 中加载迁移学习模型进行推理的正确方法是啥？

【英文标题】：What is the proper way to load a transfer learning model for inference in PyTorch?在 PyTorch 中加载迁移学习模型进行推理的正确方法是什么？

【发布时间】：2019-10-11 22:23:17

【问题描述】：

我正在使用基于 Resnet152 的迁移学习训练模型。基于 PyTorch 教程，我在保存经过训练的模型并加载它进行推理方面没有问题。但是，加载模型所需的时间很慢。我不知道我是否正确，这是我的代码：

将训练好的模型保存为状态字典：

torch.save(model.state_dict(), 'model.pkl')

加载它进行推理：

model = models.resnet152()
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(classes))
st = torch.load('model.pkl', map_location='cuda:0' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(st)
model.eval()

我对代码进行计时，发现第一行model = models.resnet152()加载的时间最长。在 CPU 上，测试一张图像需要 10 秒。所以我的想法是这可能不是加载它的正确方法？

如果我像这样保存整个模型而不是 state.dict：

torch.save(model, 'model_entire.pkl')

并像这样测试它：

model = torch.load('model_entire.pkl')
model.eval()

在同一台机器上测试一张图像只需 5 秒。

所以我的问题是：这是加载 state_dict 进行推理的正确方法吗？

【问题讨论】：

这两种方法我都用过，我建议使用后者。在前者上，您首先使用随机权重实例化模型，然后加载预训练的模型，因此它更慢是有道理的。

【参考方案1】：

在第一个代码 sn-p 中，您从 TorchVision 下载一个模型（具有随机权重），然后将您的（本地存储的）权重加载到它。

在第二个示例中，您正在加载本地存储的模型（及其权重）。

前者会更慢，因为您需要连接到托管模型的服务器并下载它，而不是本地文件，但它更可复制而不依赖于您的本地文件。此外，时间差应该是一次性初始化，并且它们应该具有相同的时间复杂度（在您执行推理时，模型已经加载到两者中，并且它们是等效的）。