Pytorch學習（三）static_dict

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

此信息與保存和加載模型和優化器相關，以供將來使用。