Pytorch之继承Module类来构造模型

简述

Module 类是nn模块里提供的一个模型构造类，是所有神经网络模块的基类，我们可以继承它来定义我们想要的模型。下面继承Module类构造多层感知机。这里定义的MLP类重载了Module类的 KaTeX parse error: Expected group after '_' at position 1: _̲_init__函数和$forward$函数。它们分别用于创建模型参数和定义前向计算。前向计算也即正向传播。

import torch
from torch import nn

class MLP(nn.Module):
    # 声明带有模型参数的层，这里声明了两个全连接层
    def __init__(self, **kwargs):
        # 调用MLP父类Module的构造函数来进行必要的初始化。这样在构造实例时还可以指定其他函数
        super(MLP, self).__init__(**kwargs)
        self.hidden = nn.Linear(784, 256) #隐藏层
        self.act = nn.ReLU()
        self.output = nn.Linear(256, 10) #输出层

    # 定义模型的前向计算，即如何根据输入x计算返回所需要的模型输出
    def forward(self, x):
        a = self.act(self.hidden(x))
        return self.output(a)

X = torch.rand(2, 784)
net = MLP()
print(net)
print(net(X))

以上的MLP类无须定义反向传播函数。系统将通过自动求梯度而自动生成反向传播所需的$backward$函数。
我们可以实例化$MLP$类得到模型变量$net$。我们初始化$net$并传入输入数据$X$做一次前向计算。其中，$net(X)$会调用$MLP$继承自$Module$类的KaTeX parse error: Expected group after '_' at position 1: _̲_call__函数,这个函数将调用$MLP$定义的$forward$函数来完成前向计算。
输出:

'''
MLP(
  (hidden): Linear(in_features=784, out_features=256, bias=True)
  (act): ReLU()
  (output): Linear(in_features=256, out_features=10, bias=True)
)
tensor([[ 0.2508,  0.0283,  0.0251,  0.1689, -0.0413,  0.2094,  0.1469, -0.0695,
          0.0132,  0.0307],
        [ 0.1728, -0.0050,  0.1031,  0.1037,  0.0778,  0.0865,  0.0227,  0.0212,
          0.2213, -0.0278]], grad_fn=<AddmmBackward0>)
'''

注意:

• 这里并没有将$Module$类命名为$Layer$(层)或者$Module$(模型)之类的名字,这是因为该类是一个可供自由组建的部件.它的子类既可以是一个层(如Pytorch提供的$Linear$类),又可以是一个模型(如这里定义的$MLP$类),或者是模型的一部分.$Module$类是一个通用的部件,事实上,Pytorch还实现了继承自$Module$的可以方面构建模型的类:如$Sequential,ModuleList,ModuleDict$等等.