PyTorch：保存权重和模型定义

Posted 2023-03-12

【中文标题】PyTorch：保存权重和模型定义

【英文标题】：PyTorch: saving both weights and model definition

【发布时间】：2021-06-13 13:44:48

【问题描述】：

在原型设计期间，我经常对 PyTorch 模型进行大量更改。例如，假设我正在试验的第一个模型是：

class Model(nn.Module):
    def __init__(self, **args):
        super().__init__()
        self.l1 = nn.Linear(128, 1)

那我再添加一层：

class Model(nn.Module):
    def __init__(self, **args):
        super().__init__()
        self.l1 = nn.Linear(128, 32)
        self.l2 = nn.Linear(32, 1)

或者可能添加一些卷积，等等。

问题在于，我执行的实验越多，我就越经常变得杂乱无章，因为我还没有找到一种直接的方法来保存模型定义及其权重，以便我可以加载以前的状态。

我知道我能做到：

torch.save('model': Model(), 'state': model.state_dict(), path)
# or directly
torch.save(model, path)

但随后加载模型还需要模型类（此处为Model）存在于当前文件中。

在Keras 你可以简单地做：

model = ...  # Get model (Sequential, Functional Model, or Model subclass)
model.save('path/to/location')

它可以保存模型的架构/配置和权重等。这意味着您可以在没有定义架构的情况下加载模型：

model = keras.models.load_model('path/to/location')

参考Keras model saving：

SavedModel 和 HDF5 文件包含：

模型的配置（拓扑） 模型的权重 模型的优化器状态（如果有）

因此，模型可以在完全相同的状态下重新实例化，无需任何 用于模型定义或训练的代码。

这就是我想在 PyTorch 中实现的目标。

PyTorch 有类似的方法吗？这些情况的最佳做法是什么？

【问题讨论】：

这能回答你的问题吗？ Best way to save a trained model in PyTorch?

抱歉在这里有点直言不讳：这些情况的最佳做法是保持井井有条。当您进行科学实验时，严谨性和组织性非常重要。您可以使用 pickle，但这仍然需要组织，因为它需要能够访问您为保存的模型定义的类/函数。

@DHudson - 这是一个有效的答案！我认为我可能忽略或滥用了 PyTorch 提供的功能（因为在 Keras 中这是可能的）。

@Datsheep - 这要求在加载之前定义 Model 类，如果可能的话，这是我想要避免的，因为原始定义可能不再可用。

您可以尝试将包含描述模型的参数的文件保存在与模型相同的目录中。这可以帮助您保持井井有条。

【参考方案1】：

由于 Pytorch 在模型中提供了极大的灵活性，因此将架构与权重一起保存在单个文件中将是一项挑战。 Keras 模型通常仅通过堆叠 keras 组件构建，但 pytorch 模型由库使用者以自己的方式编排，因此可以包含任何类型的逻辑。

我认为你有三个选择：

为您的实验提出一个有组织的架构，这样就不太可能丢失模型定义。您可以使用简单的文件，例如通过仅定义每个模型的模式命名的文件。我会推荐这种方法，因为这种组织级别可能会以其他方式使您的原型设计受益，而且开销很小。

尝试将代码与 pickle 文件一起保存。尽管有可能，但我认为这会让你陷入一个有很多潜在问题的陷阱。

使用不同的标准化方式保存模型，例如onnx。如果您不想使用选项 1，我会推荐这条路线。 Onnx 确实允许您保存 pytorch 模型的架构及其权重，但有一些缺点。例如，它只支持一些操作，因此完全自定义forward 方法或使用非矩阵操作可能不起作用。

【参考方案2】：

@D Hudson's answer 是正确的方法。但是，为了将来参考，我想添加以下对我有用的方法。

假设模型的forward 方法是固定的，即只是底层架构发生了变化，输入和输出形状相同。在这种情况下，我们只对代表整个架构的Sequential 属性感兴趣：

class Model(nn.Module):
    def __init__(self, **hparams):
        super(Model, self).__init__()
        
        # this attribute is the only thing we care about
        self.net = nn.Sequential(
            # experiment with different layers here ...
        )
        
    def forward(self, x):
        return self.net(x) # this is fixed!

然后，我们可以保存模型架构（本质上只是net 属性）及其权重：

m = Model()
# train/test/valid ...
T.save('net': m.net, 'weights': m.state_dict(), './version1.pth')

最后，加载是这样执行的：

m = Model()
checkpoint = T.load('./version1.pth')
m.net = checkpoint['net']
m.load_state_dict(checkpoint['weights'])