2.7 PyTorch中的模块简介

2.7.1 模块类

作为PyTorch深度学习框架的核心概念之一，模块在PyTorch深度学习模型的搭建过程中扮演着重要的角色。模块本身是一个类nn.Module，PyTorch的模型通过继承该类，在类的内部定义子模块的实例化，通过前向计算调用子模块，最后实现深度学习模型的搭建，具体模块类的构建如代码2.19所示。

代码2.19 PyTorch模块类的构建方法。

从代码2.19中可以看到，首先导入torch.nn库，然后基于继承nn.Module的方法构建深度学习模块。整个模块的函数主要由两部分组成：通过__init__方法初始化整个模型，forward方法对该模型进行前向计算。其中，在使用__init__方法的时候，可以在类内部初始化子模块，然后在forward方法中调用这些初始化的子模块，最后输出结果张量。在代码2.19中，把类的名字取为Model，实际上可以给类取任意的名字。因为需要调用父类nn.Module的初始化方法，这里需要使用super函数来获取当前类的父类（即nn.Module），然后调用父类的构造函数，从而初始化一些必要的变量和参数。

2.7.2 基于模块类的简单线性回归类

为了让读者更加熟悉PyTorch模块的应用，我们用PyTorch来构造一个线性回归模型，具体如代码2.20所示。

代码2.20 PyTorch线性回归模型示例。

根据前面的介绍我们已经知道，线性回归模型是输入一个特征的张量，做线性变换，输出一个预测张量。为了能够构造线性变换，我们需要知道输入特征维度大小，并且知道线性回归的权重（self.weight）和偏置（self.bias）。在forward方法中，输入一个特征张量x（大小为迷你批次大小×特征维度大小），做线性变换（使用mm方法做矩阵乘法进行线性变换），加偏置的值，最后输出一个预测的值。需要注意的是模型的初始化部分，self.weight和self.bias是模型的参数，并且一开始被初始化，使得每个分量为标准正态分布（torch.randn）。另外，需要使用nn.Parameter来包装这些参数，使之成为子模块（仅仅由参数构成的子模块），这是因为在后续训练的时候需要对参数进行优化，只有将张量转换为参数才能在后续的优化过程中被优化器访问到。

在使用线性回归模型之前，首先要做的一件事是模型的初始化，初始化的任务由初始化模型的类实例开始。对于在代码2.20中构造的线性回归模型来说，假如输入的特征大小为n，可以直接调用LinearModel（n）来构造线性回归模型的一个实例。如果需要预测m×n大小的张量对应的输出值，可以直接将m×n输入线性回归模型的实例中，如代码2.21所示。

代码2.21 PyTorch线性回归模型调用方法实例。

2.7.3 线性回归类的实例化和方法调用

对于PyTorch的模块，有一些常用的方法可以在训练和预测的时候调用。

1．使用named_parameters方法和parameters方法获取模型的参数

通过调用named_parameters方法，返回的是Python的一个生成器（Generator），通过访问生成器的对象得到的是该模型所有参数的名称和对应的张量值。通过调用parameters方法，返回的也是一个生成器，访问生成器的结果是该模型的所有参数对应张量的值。PyTorch的优化器直接接受模型的参数生成器作为函数的参数，并且会根据梯度来优化生成器里的所有张量（需要调用反向传播函数）。

2．使用train和eval方法进行模型训练和测试状态的转换

在模型的使用过程中，有些子模块（如丢弃层和批次归一化层等）有两种状态，即训练状态和预测状态，PyTorch的模型经常需要在两种状态中相互转换。通过调用train方法会把模块（包括所有的子模块）转换到训练状态，调用eval方法会把模块（包括所有的子模块）转换到预测状态。PyTorch的模型在不同状态下的预测准确率会有差异，在训练模型的时候需要转换为训练状态，在预测的时候需要转换为预测状态，否则最后模型预测准确率可能会降低，甚至会得到错误的结果。

3．使用named_buffers方法和buffers方法获取张量的缓存

除通过反向传播得到梯度来进行训练的参数外，还有一些参数并不参与梯度传播，但是会在训练中得到更新，这种参数称为缓存（Buffer），其中具体的例子包括批次归一化层的平均值（Mean）和方差（Variance）。通过在模块中调用register_buffer方法可以在模块中加入这种类型的张量，通过named_buffers可以获得缓存的名字和缓存张量的值组成的生成器，通过buffers方法可以获取缓存张量值组成的生成器。

4．使用named_children方法和children方法获取模型的子模块

有时需要对模块的子模块进行迭代，这时就需要使用named_children方法和children方法来获取子模块名字、子模块的生成器，以及只有子模块的生成器。由于PyTorch模块的构造可以嵌套，所以子模块还有可能有自身的子模块，如果要获取模块中所有模块的信息，可以使用named_modules和modules来（递归地）得到相关信息。

5．使用apply方法递归地对子模块进行函数应用

如果需要对PyTorch所有的模块应用一个函数，可以使用apply方法，通过传入一个函数或者匿名函数（通过Python语言的lambda关键字定义）来递归地应用这些函数。传入的函数以模块作为参数，在函数内部对模块进行修改。

6．改变模块参数数据类型和存储的位置

除对模块进行修改外，在深度学习模型的构建中还可能对参数进行修改。和张量的运算一样，可以改变模块的参数所在的设备（CPU或者GPU），具体可以通过调用模块自带的cpu方法和cuda方法来实现。另外，如果需要改变参数的数据类型，可以通过调用to方法加上需要转变的目标数据类型来实现（也可以使用具体的一些方法，比如float方法会转换所有的参数为单精度浮点数，half方法会转换所有的参数为半精度浮点数，double方法会转换所有的参数为双精度浮点数）。具体调用的实例可以参考代码2.22。

代码2.22 PyTorch模块方法调用实例。