WHAT IS A STATE_DICT IN PYTORCH

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

 

Introduction

 

如果有兴趣从PyTorch保存或加载模型，则state_dict是不可或缺的实体。由于state_dict对象是Python字典，因此可以轻松地保存，更新，更改和还原它们，从而为PyTorch模型和优化器增加了很多模块化。请注意，只有具有可学习参数的层（卷积层，线性层等）和已注册的缓冲区（batchnorm的running_mean）才在模型的中具有条目state_dict。优化器对象（torch.optim）也有一个state_dict，其中包含有关优化器状态以及所用超参数的信息。在文中，我们将看到如何state_dict与简单模型一起使用。

 

Setup

在开始之前，需要先安装（如果torch尚不可用）​

 

pip install torch

 

Steps

1. 导入所有必需的库以加载我们的数据
2. 定义和初始化神经网络
3. 初始化优化器
4. 访问模型和优化器 state_dict

 

1. Import necessary libraries for loading our data

我们将使用torch及其子公司torch.nn 和torch.optim。

import torchimport torch.nn as nnimport 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 xnet = Net()print(net)

 

3. Initialize the optimizer

我们将要用带有momentum的SGD

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])

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

​

接下来，给大家介绍一下租用GPU做实验的方法，我们是在智星云租用的GPU，使用体验很好。具体大家可以参考：智星云官网： http://www.ai-galaxy.cn/，淘宝店：https://shop36573300.taobao.com/公众号: 智星AI