PyTorch基础教程学习笔记（六）： nn.Module

前言：前面介绍了如何自定义一个模型——通过继承nn.Module类来实现，在__init__构造函数中申明各个层的定义，在forward中实现层之间的连接关系，实际上就是前向传播的过程。

事实上，在pytorch里面自定义层也是通过继承自nn.Module类来实现的，pytorch里面一般是没有层的概念，层也是当成一个模型来处理的，这里和keras是不一样的。当然也可以直接通过继承torch.autograd.Function类来自定义一个层，但是这很不推荐，不提倡，至于为什么后面会介绍。记住一句话，keras更加注重的是层Layer、pytorch更加注重的是模型Module。

所以本文就专门来介绍如何通过nn.Module类来实现自定义层。

从系统预定义的层说起

Linear层的代码

import math
import torch
from torch.nn.parameter import Parameter
from .. import functional as F
from .. import init
from .module import Module
from ..._jit_internal import weak_module, weak_script_method
 
class Linear(Module):
    __constants__ = ['bias']
 
    def __init__(self, in_features, out_features, bias=True):
        super(Linear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.Tensor(out_features, in_features))
        if bias:
            self.bias = Parameter(torch.Tensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()
 
    def reset_parameters(self):
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)
 
    @weak_script_method
    def forward(self, input):
        return F.linear(input, self.weight, self.bias)
 
    def extra_repr(self):
        return 'in_features={}, out_features={}, bias={}'.format(
            self.in_features, self.out_features, self.bias is not None
        )

Conv2d类的实现

class Conv2d(_ConvNd):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1,
                 bias=True, padding_mode='zeros'):
        kernel_size = _pair(kernel_size)
        stride = _pair(stride)
        padding = _pair(padding)
        dilation = _pair(dilation)
        super(Conv2d, self).__init__(
            in_channels, out_channels, kernel_size, stride, padding, dilation,
            False, _pair(0), groups, bias, padding_mode)
 
    @weak_script_method
    def forward(self, input):
        if self.padding_mode == 'circular':
            expanded_padding = ((self.padding[1] + 1) // 2, self.padding[1] // 2,
                                (self.padding[0] + 1) // 2, self.padding[0] // 2)
            return F.conv2d(F.pad(input, expanded_padding, mode='circular'),
                            self.weight, self.bias, self.stride,
                            _pair(0), self.dilation, self.groups)
        return F.conv2d(input, self.weight, self.bias, self.stride,
                        self.padding, self.dilation, self.groups)

初步总结

在前面的文章里面说过，torch里面实现神经网络有两种方式

（1）高层API方法：使用torch.nn.****来实现；

（2）低层API方法：使用低层函数方法，torch.nn.functional.****来实现；

其中，我们推荐使用高层API的方法，原因如下：

高层API是使用类的形式来包装的，既然是类就可以存储参数，比如全连接层的权值矩阵、偏置矩阵等都可以作为类的属性存储着，但是低层API仅仅是实现函数的运算功能，没办法保存这些信息，会丢失参数信息，但是高层API是依赖于低层API的计算函数的，比如上面的两个层：

• Linear高级层——>低层F.linear() 函数
• Conv2d高级层——>低层F.conv2d()函数

---

自定义层的步骤

要实现一个自定义层大致分以下几个主要的步骤：

1. 自定义一个类，继承自Module类，并且一定要实现两个基本的函数，第一是构造函数__init__，第二个是层的逻辑运算函数，即所谓的前向计算函数forward函数。
2. 在构造函数_init__中实现层的参数定义。比如Linear层的权重和偏置，Conv2d层的in_channels, out_channels, kernel_size, stride=1,padding=0, dilation=1, groups=1,bias=True, padding_mode='zeros'这一系列参数；
3. 在前向传播forward函数里面实现前向运算。这一般都是通过torch.nn.functional.***函数来实现，当然很多时候我们也需要自定义自己的运算方式。如果该层含有权重，那么权重必须是nn.Parameter类型，关于Tensor和Variable（0.3版本之前）与Parameter的区别请参阅相关的文档。简单说就是Parameter默认需要求导，其他两个类型则不会。另外一般情况下，可能的话，为自己定义的新层提供默认的参数初始化，以防使用过程中忘记初始化操作。
4. 补充：一般情况下，我们定义的参数是可以求导的，但是自定义操作如不可导，需要实现backward函数。

Tensor

       pytorch中的Tensor类似于numpy中的array，而直接用tensor的原因，是因为tensor能够更方便地在GPU上进行运算。pytorch为tensor设计了许多方便的操作，同时tensor也可以轻松地和numpy数组进行相互转换。

Variable

   Variable是对Tensor的封装，是Tensor的一个Wrapper，其中保存了Variable的创造者，Variable的值（tensor），还有Variable的梯度(Variable)，即每一个Variable被构建的时候，都包含三个属性：

• Variable中所包含的tensor
• tensor的梯度 .grad
• 以何种方式得到这种梯度 .grad_fn

操作与tensor基本一致。 之所以有Variable这个数据结构，是为了引入计算图（自动求导），方便构建神经网络。 Variable的前向过程的计算包括两个部分的计算，一个是其值的计算（即，Tensor的计算），还有就是Variable标签的计算。标签指的是什么呢？如果您看过PyTorch的官方文档 Excluding subgraphs from backward 部分的话，您就会发现Variable还有两个标签：requires_grad和volatile。标签的计算指的就是这个。      

        简单举个例子：

from torch.autograd import Variable
a = torch.randn(10, 5)
b = torch.randn(10, 5)
x = Variable(a, requires_grad=True)
y = Variable(b, requires_grad=True)
z = x + y
z.backward()
x.grad            # x的梯度 10x1 的全1 tensor
z.grad_fn         #  

       通过调用backward()，我们可以对某个Variable（譬如说y）进行 一次自动求导，但如果我们再对这个Variable进行一次backward()操作，会发现程序报错。这是因为PyTorch默认做完一次自动求导后，就把计算图丢弃了。我们可以通过设置retain_graph来实现多次求导。

Parameter

       我们知道网络中存在很多参数，这些参数需要在网络训练的过程中实时更新（一个batch更新一次），完成“学习”的过程，譬如最直观的梯度下降法更新参数w：

w.data = w.data - lr * w.grad.data  # lr 是学习率

• 网络中若是有100个参数，都要手写更新代码吗？1000个呢？10000个呢......
• Variable默认是不需要求梯度的，那还需要手动设置参数 requires_grad=True
• Variable因为要多次反向传播，那么在bcakward的时候还要手动注明参数w.backward(retain_graph=True)

​​​​        Pytorch主要通过引入nn.Parameter类型的变量和optimizer机制来解决了这个问题。Parameter是Variable的子类，本质上和后者一样，只不过parameter默认是求梯度的，同时一个网络net中的parameter变量是可以通过 net.parameters() 来很方便地访问到的，只需将网络中所有需要训练更新的参数定义为Parameter类型，再佐以optimizer，就能够完成所有参数的更新了，具体如下：

  class Net(Module):
        def __init__(self, a, b, ...):
                super(net, self).__init__()
                self...   #  parameters
                self...    # layers
        def forward(self):
                x = ...
                x = ...    # 数据流
                return x
net = Net(a, b, ...)
net.train()
...
optimizer = torch.optim.SGD(net.parameters(), lr=1e-1)
# 然后在每一个batch中，调用optimizer.step()即可完成参数更新了（loss.backward()之后）

---

自定义层的简单例子

比如要实现一个简单的层，这个层的功能是

即输入X的平方再加上一个偏置项，再开跟根号，然后再乘以权值矩阵w，那要怎么做呢，按照上面的定义过程，我们先定义一个这样的层（即一个类），代码如下：

定义一个自定义层MyLayer

# 定义一个 my_layer.py
import torch
 
class MyLayer(torch.nn.Module):
    '''
    因为这个层实现的功能是：y=weights*sqrt(x2+bias),所以有两个参数：
    权值矩阵weights
    偏置矩阵bias
    输入 x 的维度是（in_features,)
    输出 y 的维度是（out_features,) 故而
    bias 的维度是（in_fearures,)，注意这里为什么是in_features,而不是out_features，注意体会这里和Linear层的区别所在
    weights 的维度是（in_features, out_features）注意这里为什么是（in_features, out_features）,而不是（out_features, in_features），注意体会这里和Linear层的区别所在
    '''
    def __init__(self, in_features, out_features, bias=True):
        super(MyLayer, self).__init__()  # 和自定义模型一样，第一句话就是调用父类的构造函数
        self.in_features = in_features
        self.out_features = out_features
        self.weight = torch.nn.Parameter(torch.Tensor(in_features, out_features)) # 由于weights是可以训练的，所以使用Parameter来定义
        if bias:
            self.bias = torch.nn.Parameter(torch.Tensor(in_features))             # 由于bias是可以训练的，所以使用Parameter来定义
        else:
            self.register_parameter('bias', None)
 
    def forward(self, input):
        input_=torch.pow(input,2)+self.bias
        y=torch.matmul(torch.sqrt(input_),self.weight)
        return y

自定义模型并且训练

import torch
from my_layer import MyLayer # 自定义层
 
 
N, D_in, D_out = 10, 5, 3  # 一共10组样本，输入特征为5，输出特征为3 
 
# 先定义一个模型
class MyNet(torch.nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()  # 第一句话，调用父类的构造函数
        self.mylayer1 = MyLayer(D_in,D_out)
 
    def forward(self, x):
        x = self.mylayer1(x)
 
        return x
 
model = MyNet()
print(model)
'''运行结果为：
MyNet(
  (mylayer1): MyLayer()   # 这就是自己定义的一个层
)
'''

下面开始训练

# 创建输入、输出数据
x = torch.randn(N, D_in)  #（10，5）
y = torch.randn(N, D_out) #（10，3）
 
 
#定义损失函数
loss_fn = torch.nn.MSELoss(reduction='sum')
 
learning_rate = 1e-4
#构造一个optimizer对象
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
 
for t in range(10): # 
    
    # 第一步：数据的前向传播，计算预测值p_pred
    y_pred = model(x)
 
    # 第二步：计算计算预测值p_pred与真实值的误差
    loss = loss_fn(y_pred, y)
    print(f"第 {t} 个epoch, 损失是 {loss.item()}")
 
    # 在反向传播之前，将模型的梯度归零，这
    optimizer.zero_grad()
 
    # 第三步：反向传播误差
    loss.backward()
 
    # 直接通过梯度一步到位，更新完整个网络的训练参数
    optimizer.step()

那么调用forward方法的具体流程是什么样的呢？具体流程是这样的：

以一个Module为例：
1. 调用module的call方法
2. module的call里面调用module的forward方法
3. forward里面如果碰到Module的子类，回到第1步，如果碰到的是Function的子类，继续往下
4. 调用Function的call方法
5. Function的call方法调用了Function的forward方法。
6. Function的forward返回值
7. module的forward返回值
8. 在module的call进行forward_hook操作，然后返回值

上述中“调用module的call方法”是指nn.Module 的__call__方法。定义__call__方法的类可以当作函数调用，具体参考Python的面向对象编程。也就是说，当把定义的网络模型model当作函数调用的时候就自动调用定义的网络模型的forward方法。

 程序的运行结果为：

第 0 个epoch, 损失是 29.47430419921875
第 1 个epoch, 损失是 29.46476936340332
第 2 个epoch, 损失是 29.455242156982422
第 3 个epoch, 损失是 29.445714950561523
第 4 个epoch, 损失是 29.436201095581055
第 5 个epoch, 损失是 29.426692962646484
第 6 个epoch, 损失是 29.417190551757812
第 7 个epoch, 损失是 29.40769386291504
第 8 个epoch, 损失是 29.398204803466797
第 9 个epoch, 损失是 29.38872528076172

注意：sqrt()的存在导致loss可能为nan

sqrt(x) 函数的定义域为 [0, 无穷大)
sqrt(x) 的导函数的定义域 却是 (0, 无穷大)

这些函数定义域跟导函数的定义域不一样，正向传播可以得到正常结果，但是一旦backward就会得到Nan

如何解决
让输入的值符合sqrt的导函数定义域就可以解决该问题了。举个例子：设 x 的定义域为 [0, 无穷大) ，给 x 加个很小的数，例如1e-8，使其输入值的定义域略微往右偏移，就可以避开 0 这个未定义值了；y = sqrt(x + 1e-8)