【pytorch】pytorch中的state_dict解读

在pytorch中，torch.nn.Module模块中的state_dict变量存放训练过程中需要学习的权重和偏执系数，state_dict作为python的字典对象将每一层的参数映射成tensor张量。

import torch.nn as nn
class A(nn.Module):
    def __init__(self):
        super(A,self).__init__()
        self.l1=nn.Linear(2,2)

    def forward(self,x):
        return self.l1(x)

class B(A):
    def __init__(self):
        super(B,self).__init__()
        self.a=A()
        self.l2=nn.Linear(2,1)
    def forward(self,x):
        x=self.a(x)
        return self.l2(x)

b=B()

dict1=b.state_dict()
for e in b.state_dict():
    print(e,'\t',dict1[e])

输出：

l1.weight      tensor([[ 0.4777,  0.4941],
        [ 0.0633, -0.2611]])
l1.bias      tensor([ 0.3897, -0.3371])
a.l1.weight      tensor([[-0.6648, -0.1620],
        [ 0.4378,  0.6222]])
a.l1.bias      tensor([-0.6461, -0.0242])
l2.weight      tensor([[-0.2424,  0.4311]])
l2.bias      tensor([-0.4948])