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

这取决于您以后要对模型执行的操作。

出于恢复培训的目的

如果您想保存模型以便以后进行训练，则除了模型本身的state_dict以外，还需要更多。您还需要保存优化器的状态，纪元，得分等。

state = {
    'epoch': epoch,
    'state_dict': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    ...
}
torch.save(state, filepath)

# load
model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])

出于推理目的

仅模型的state_dict就足够了。但是请确保在加载模型后调用eval模式，以便batchnorm或dropout层将在评估模式下而不是训练模式下工作。

torch.save(model.state_dict(), filepath)

# load
model.load_state_dict(torch.load(filepath))
model.eval()

如果需要，请参考官方best practices guide。

Answer 2

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

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

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

在PyTorch中加载用于推理的传输学习模型的正确方法是什么？

2 个答案: