pytorch(9)： nn.Module深入理解

测试版本:
torch1.7.1 + CPU

python 3.6

nn.Module的构造函数的源代码：

def __init__(self):
    self._parameters = OrderedDict()
    self._modules = OrderedDict()
    self._buffers = OrderedDict()
    self._backward_hooks = OrderedDict()
    self._forward_hooks = OrderedDict()
    self.training = True

• parameters: 保存用户设置的参数。但不保存子module的参数
• modules:保存子module
• buffers:缓存。有时前向传播时需要使用到上一次前向传播的结果
• hooks:钩子技术
• training: 标识训练阶段与测试阶段。(如DropOut层在两阶段的运行状态不同，所以需要标记)

1、添加parameters和子module的测试:

from torch import nn
import torch as t
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.param1 = nn.Parameter(t.rand(3,3))
        #等同于:self.register_parameter('param1',nn.Parameter(t.rand(3,3)))

        self.subModule1 = nn.Linear(3,4)
    def forward(self,input):
        x = self.param1.mm(input)
        x = self.subModule1(x)
        return x

net = Net()
net

Net(
  (subModule1): Linear(in_features=3, out_features=4, bias=True)
)

net._parameters

OrderedDict([('param1',
              Parameter containing:
              tensor([[0.4540, 0.3877, 0.9361],
                      [0.2931, 0.4030, 0.7346],
                      [0.8997, 0.4647, 0.3729]], requires_grad=True))])

net._modules

OrderedDict([('subModule1', Linear(in_features=3, out_features=4, bias=True))])

2、保存/加载 module

t.save(net.state_dict(),'net.path')#这时候会创建一个net.path文件

net2 = Net()
net2.load_state_dict(t.load('net.path'))
net2._parameters

OrderedDict([('param1',
              Parameter containing:
              tensor([[0.4540, 0.3877, 0.9361],
                      [0.2931, 0.4030, 0.7346],
                      [0.8997, 0.4647, 0.3729]], requires_grad=True))])