PyTorch中的state_dict是什么

Created On: Apr 17, 2020 | Last Updated: Feb 06, 2024 | Last Verified: Nov 05, 2024

在 PyTorch 中，torch.nn.Module 模型的可学习参数（即权重和偏置）包含在模型的参数中（通过 model.parameters() 访问）。state_dict 只是一个 Python 字典对象，它将每层映射到其参数张量。

简介

如果您对从 PyTorch 保存或加载模型感兴趣，state_dict 是一个不可或缺的实体。由于 state_dict 对象是 Python 字典，它们可以轻松保存、更新、更改和恢复，从而为 PyTorch 模型和优化器添加了很大的模块化。请注意，只有具有可学习参数的层（卷积层、线性层等）和注册的缓冲区（如 batchnorm 的运行平均值）在模型的 state_dict 中有条目。优化器对象（torch.optim）也有一个 state_dict，其中包含有关优化器状态的信息以及使用的超参数。在本教程中，我们将了解如何使用简单的模型来使用 state_dict。

环境设置

在开始之前，我们需要安装``torch``，如果尚未安装。

pip install torch

步骤

1. 导入加载数据所需的所有库。

2. 定义并初始化神经网络。

3. 初始化优化器

4. 访问模型和优化器的 state_dict

1. 导入加载数据所需的库。

在本教程中，我们将使用``torch``及其组件``torch.nn``和``torch.optim``。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

2. 定义并初始化神经网络。

为了举例说明，我们将创建一个用于训练图像的神经网络。欲了解更多信息，请参阅定义神经网络教程。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()
print(net)

3. 初始化优化器

我们将使用带动量的 SGD。

optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

4. 访问模型和优化器的 state_dict

现在我们已经构建了我们的模型和优化器，我们可以了解它们各自的 state_dict 属性中保存了什么内容。

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in net.state_dict():
    print(param_tensor, "\t", net.state_dict()[param_tensor].size())

print()

# Print optimizer's state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

此信息与保存和加载模型及优化器以供将来使用有关。

恭喜您！您已成功在 PyTorch 中使用了 state_dict。

了解更多

查看这些其他教程以继续学习：

• PyTorch 中用于推理的模型保存和加载

• 在 PyTorch 中保存和加载通用检查点