state_dict字典状态

state_dict字典状态

state_dict字典状态本质就是一个Python字典，键值对是每个网络层和其对应的参数张量。

state-dict包含带有学习参数的网络层和注册的缓存

优化器对象包含一个状态字典，包含优化器状态的信息和使用的超参数。

import torch.nn as nn
import torch.optim as optim

class ModelClass(nn.Module):
    def __init__(self):
        super(ModelClass,self).__init__()
        self.convolution_1 = nn.Conv2d(3,6,5)
        self.pool = nn.MaxPool2d(2,2)
        self.convolution_2 = nn.Conv2d(6,16,5)
        self.full_connection_1 = nn.Linear(16*5*5,120)
        self.full_connection_2 = nn.Linear(120,84)
        self.full_connection_3 = nn.Linear(84,10)

    def forward(self,x):
        x = self.pool(F.relu(self.convolution_1(x)))
        x = self.pool(F.relu(self.convolution_2(x)))
        x = x.view(-1,16*5*5)
        x = F.relu(self.full_connection_1(x))
        x = F.relu(self.full_connection_2(x))
        x = self.full_connection_3(x)
        return x

# Initialize model
model = ModelClass()

#Initialize optimizer
optimizer = optim.SGD(model.parameters(),lr = 0.001,momentum = 0.9)

for paramter_tensor in model.state_dict():
    print(paramter_tensor,"\t",model.state_dict()[paramter_tensor].size)

输出：