在 Pytorch 中节省一层的重量

Question

我想保存一个模型的重量，而不是像这样保存整个模型的重量：

torch.save(model, 'model.pth')

而是，只有一层。例如，假设我定义了这样的一层：

self.conv_up3 = convrelu(256 + 512, 512, 3, 1)

如何只保存这一层的权重。还有我如何为这一层加载它。

Answer 1

您可以对特定层的save/get参数执行以下操作：

specific_params = self.conv_up3.state_dict()
# save/manipulate `specific_params` as you want

同样，将参数加载到该特定层：

self.conv_up3.load_state_dict(params)

您可以这样做，因为每一层本身就是一个神经网络（nn.Module 实例）。