注意
跳转到结尾 下载完整的示例代码。
PyTorch 中的 state_dict 是什么#
创建日期：2020 年 4 月 17 日 | 最后更新：2024 年 2 月 6 日 | 最后验证：2024 年 11 月 5 日
在 PyTorch 中，torch.nn.Module 模型的学习参数（即权重和偏置）包含在模型的参数中（通过 model.parameters() 访问）。state_dict 只是一个 Python 字典对象，它将每个层映射到其参数张量。
简介#
如果您对保存或加载 PyTorch 中的模型感兴趣，state_dict 是一个不可或缺的实体。由于 state_dict 对象是 Python 字典，因此可以轻松保存、更新、修改和恢复它们，为 PyTorch 模型和优化器增加极大的模块化程度。请注意，只有具有可学习参数的层（卷积层、线性层等）和注册的缓冲区（batchnorm 的 running_mean）才具有模型 state_dict 中的条目。优化器对象 (torch.optim) 也具有 state_dict，其中包含有关优化器状态以及所用超参数的信息。在本教程中，我们将了解如何将 state_dict 与一个简单的模型一起使用。
设置#
在开始之前，如果尚未安装 torch，我们需要安装它。
pip install torch
步骤#
导入加载数据所需的所有库
定义并初始化神经网络
初始化优化器
访问模型和优化器的
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。
了解更多#
查看这些其他秘籍以继续您的学习
