如何加载和使用PyTorch（.pth.tar）模型

Question

我对Torch不太熟悉，我主要使用Tensorflow。但是，我需要使用在火炬中重新训练的再训练初始模型。由于为我的特定应用程序重新训练初始模型所需的大量计算资源，我想使用已经重新训练的模型。

此模型保存为.pth.tar文件。

我希望能够首先加载这个模型。到目前为止，我已经能够弄清楚我必须使用以下内容：

model = torch.load('iNat_2018_InceptionV3.pth.tar', map_location='cpu')

这似乎有效，因为print(model)打印出一大组数字和其他值，我认为这是权重的偏差值。

在此之后，我需要能够使用它对图像进行分类。我无法弄清楚这一点。我该如何格式化图像？图像应该转换成数组吗？在此之后，我如何将输入数据传递到网络？

Answer 1

你基本上需要和tensorflow一样。也就是说，当您存储网络时，只会存储参数（即网络中的可训练对象），而不会存储“胶水”，这就是您使用训练模型所需的所有逻辑。因此，如果您有.pth.tar文件，则可以加载它，从而覆盖已定义模型的参数值。

这意味着保存/加载模型的一般过程如下：

• 写你的网络定义（即你的nn.Module对象）
• 以您想要的方式训练或以其他方式更改网络参数
• 使用torch.save保存参数
• 当你想使用那个网络时，使用相同的nn.Module对象定义来首先实例化一个pytorch网络
• 然后使用torch.load覆盖网络参数的值

以下是关于如何执行此操作的一些参考资料的讨论：pytorch forums

这是一个超短的mwe：

# to store
torch.save({
    'state_dict': model.state_dict(),
    'optimizer' : optimizer.state_dict(),
}, 'filename.pth.tar')

# to load
checkpoint = torch.load('filename.pth.tar')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])