python - 在 PyTorch 中加载用于推理的迁移学习模型的正确方法是什么？

Question

我正在使用基于 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 进行推理的正确方法吗？

Answer 1

在第一个代码片段中，您从 TorchVision 下载模型(具有随机权重)，然后将您的(本地存储的)权重加载到它。
在第二个示例中，您正在加载本地存储的模型(及其权重)。
前者会更慢，因为您需要连接到托管模型的服务器并下载它，而不是本地文件，但它更易于复制，不依赖于您的本地文件。此外，时间差应该是一次性初始化，并且它们应该具有相同的时间复杂度(因为在您执行推理时，模型已经加载到两者中，并且它们是等效的)。