pytorch中模型和參數的存儲

 

在PyTorch中， torch.nn.Module 模型的可學習參數（即權重和偏差）包含在模型的參數中，
（使用 c可以進行訪問）。
state_dict 是Python字典對象，它將每一層映射到其參數張量。注意，只有具有可學習參數的層（如卷積層，線性層等）的模型才具有 state_dict 這一項。目標優化 torch.optim 也有 state_dict 屬性，它包含有關優化器的狀態信息，以及使用的超參數。
因爲state_dict的對象是Python字典，所以它們可以很容易的保存、更新、修改和恢復，爲PyTorch模型和優化器添加了大量模塊。
下面通過從簡單模型訓練一個分類器中來了解一下 state_dict 的使用

 

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
class theModelClass(nn.Module):
    def __init__(self):
        super(theModelClass, 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


model = theModelClass()

# 初始化 優化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 打印模型的狀態字典
print("Model's state_dict:" )
for param_tensor in model.state_dict():
    print(param_tensor, '\t', model.state_dict()[param_tensor].size())

# 打印優化器的轉檯字典
print(" Optimizer's state_dict; ")
for var_name in optimizer.state_dict():
    print(var_name, '\t', optimizer.state_dict()[var_name])

# 輸出

"""
Model's state_dict:
conv1.weight torch.Size([6, 3, 5, 5])
conv1.bias torch.Size([6])

conv2.weight torch.Size([16, 6, 5, 5])
conv2.bias torch.Size([16])

fc1.weight torch.Size([120, 400])  
fc1.bias torch.Size([120])

fc2.weight torch.Size([84, 120])
fc2.bias torch.Size([84])

fc3.weight torch.Size([10, 84])
fc3.bias torch.Size([10])

Optimizer's state_dict:
state {}
param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay':
0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072,
4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

"""