在PyTorch中,一个torch.nn.Module可学习参数(即权重和偏差)。模型包含在model的参数中(使用model.parameters()访问)。state_dict只是一个Python字典对象,它将每一层映射到其参数张量。
介绍
如果您对从PyTorch保存或加载模型感兴趣,那么state_dict是一个完整的实体。因为state_dict对象是Python字典,所以可以很容易地保存、更新、修改和恢复它们,从而为PyTorch模型和优化器增加了大量的模块化。注意,只有具有可学习参数的层(卷积层、线性层等)和注册缓冲区(batchnorm s running_mean)在模型s state_dict中有条目。优化器对象(电筒.optim)也有一个state_dict,它包含关于优化器状态和使用的超参数的信息。在这个食谱中,我们将看到state_dict是如何与一个简单的模型一起使用的。
步骤
1. 导入包
2. 定义和初始化神经网络
3. 初始化优化器
4. 访问模型和优化器 state_dict
1. Import necessary libraries for loading our data
import torch
import torch.nn as nn
import torch.optim as optim
2. Define and intialize the neural network
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. Initialize the optimizer
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
4. Access the model and optimizer 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])
此信息与保存和加载模型和优化器相关,以供将来使用。