基于PyTorch学AI——构建网络模型

有了前面的铺垫，终于可以利用PyTorch构建一个神经网络了。
在这里插入图片描述

神经网络由很多节点组成，每个节点有输入输出，所有节点组成一个网络结构，整体网络也有输入和输出，通过输入训练数据，输出标签，完成一个特定任务。
PyTorch中利用nn.Module类实现节点功能。

构建神经网络

先说说基本逻辑，神经网络由层和模块组成，可以对数据执行某种操作。
torch.nn命名空间提供了构建一个神经网络所需的各种模块，每个模块都继承自nn.Module，一个模块可以包含多个其它模块，利用这种嵌套结构可以轻松构建出复杂的神经网络模型。

下面就开始构建FashionMNIST数据集的图片分类模型。
先导入相关库，最重要的是nn库。

import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

指定训练设备

如果有GPU，尽量使用GPU训练模型。
利用torch.cuda.is_available()方法检测cuda是否可用，如果不可用，则使用CPU。

device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

定义模型网络

自定义的神经网络模型都继承自nn.Module，需要实现init方法和forward方法。
__init__方法中定义其子模块，构成整个网络。
forward方法，实现前向运算逻辑，不需要手动调用。

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

定义一个NeuralNetwork模块，有多个子模块组成。

• flatten 子模块
  用于把一个tensor”压扁“。看例子秒懂：

>>> input = torch.randn(32, 1, 5, 5)
>>> # With default parameters
>>> m = nn.Flatten()
>>> output = m(input)
>>> output.size()
torch.Size([32, 25])
>>> # With non-default parameters
>>> m = nn.Flatten(0, 2)
>>> output = m(input)
>>> output.size()
torch.Size([160, 5])

作为模型的第一个子模块，用于把输入的数据从多维tensor”压扁“成2维tensor。
在这个例子中，Flatten层会把每一个28 * 28 图片样本转化为784个像素值的一维tensor。

• Sequential 容器
  Sequential模块是一个顺序容器，会自动把其包含的多个子模块按顺序执行，子模块之间首尾链接，上一个模块的输出作为下一个模块的输入。
  上面定义中包括5个子模块，3个Linear模块和2个ReLU模块。

• Linear 模块
  Linear模块定义了线性变换，有两个参数，分别是输入特征的维度和输出特征（隐藏层）的维度。
  还有个参数bias取的默认值true，代表是否有偏置项。

• ReLU模块
  ReLU是激活函数层，为整个模型增加非线性元素。
  这个资料太多了，感兴趣的自查，不赘述。

整个模型定义最后一个Linear的输出就是整个模型的输出，在这个场景下输出的就是10个类型的概率值。
下面代码会创建一个NeuralNetwork类的实例并移入device，并打印模型结构。

model = NeuralNetwork().to(device)
print(model)

输出结果：

通过上图，很直观的看到整个模型的数据流转情况，包括每一个子模块以及输入输出。

把训练数据传给模型实例，会自动调用model的forward()方法，不需要手动调用。
例如如下代码：

X = torch.rand(1, 28, 28, device=device)
logits = model(X)
print(f"logits: {logits}")
pred_probab = nn.Softmax(dim=1)(logits)
print(f"pred_probab: {pred_probab}")
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

输入一条[1 , 28 , 28]的rand数据模拟样本数据。
首先，Flatten子模块会把样本结构变化为[1,784]的tensor。
然后进入linear_relu_stack，经过一系列子模块的运算，输出[1,10]的tensor。
模型返回一个[1,10]tensor，是样本在10个分类的output，通过nn.Softmax函数得到每个分类的概率，最后通过argmax方法返回概率最大的分类索引。

上面代码把每个步骤的输出都打印下来，结果如下：

nn.Sequential

看源码：

    def __init__(self, *args):
        super().__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            for key, module in args[0].items():
                self.add_module(key, module)
        else:
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)

循环调用add_module方法，把所有子模块作为参数传给方法。

 def add_module(self, name: str, module: Optional['Module']) -> None:
        r"""Add a child module to the current module.

        The module can be accessed as an attribute using the given name.

        Args:
            name (str): name of the child module. The child module can be
                accessed from this module using the given name
            module (Module): child module to be added to the module.
        """
        if not isinstance(module, Module) and module is not None:
            raise TypeError(f"{torch.typename(module)} is not a Module subclass")
        elif not isinstance(name, str):
            raise TypeError(f"module name should be a string. Got {torch.typename(name)}")
        elif hasattr(self, name) and name not in self._modules:
            raise KeyError(f"attribute '{name}' already exists")
        elif '.' in name:
            raise KeyError(f"module name can't contain \".\", got: {name}")
        elif name == '':
            raise KeyError("module name can't be empty string \"\"")
        for hook in _global_module_registration_hooks.values():
            output = hook(self, name, module)
            if output is not None:
                module = output
        self._modules[name] = module

前面代码完成数据校验，最重要的代码在最后，相当于把所有子模块放到_modules: Dict[str, Optional[‘Module’]]对象。

再看forward方法：

    def forward(self, input):
        for module in self:
            input = module(input)
        return input

逻辑很简单，循环依次调用所有模块，每次的输出作为下一个模块的输入。
这里直接用module(input)，相当于调用模块的forward方法，这里的细节，以后通过分析nn.Module的代码详细了解。

nn.Flatten

详细看下Flatten的源码：

    def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None:
        super().__init__()
        self.start_dim = start_dim
        self.end_dim = end_dim

    def forward(self, input: Tensor) -> Tensor:
        return input.flatten(self.start_dim, self.end_dim)

通过tensor的flatten方法实现功能，传参默认是1和-1，相当于把输入的第2个维度到最后一个维度压平。

nn.ReLU

实现ReLU激活函数功能。
源码：

    def __init__(self, inplace: bool = False):
        super().__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.relu(input, inplace=self.inplace)

调用F的relu方法。其中，F是torch.nn.functional的别名。

nn.Softmax

再看一下Softmax模块。
神经网络的最后一个线性层返回logits，传递给nn.Softmax模块。
logits被缩放[0，1]区间，表示每个类别的预测概率值。
dim参数表示要处理的维度。

内部实现也很简单：

class Softmax(Module):
    def __init__(self, dim: Optional[int] = None) -> None:
        super().__init__()
        self.dim = dim

    def forward(self, input: Tensor) -> Tensor:
        return F.softmax(input, self.dim, _stacklevel=5)

具体实现在F中，基本就是数学公式的实现。

打印参数

可以通过model的 parameters() 或 named_parameters() 方法打印模型的所有参数。

for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.size()} | Values : {param[:2]} \n")

输出：

named_parameters方法定义在Module类中，所有子类都可以调用。

    def named_parameters(
            self,
            prefix: str = '',
            recurse: bool = True,
            remove_duplicate: bool = True
    ) -> Iterator[Tuple[str, Parameter]]:
        gen = self._named_members(
            lambda module: module._parameters.items(),
            prefix=prefix, recurse=recurse, remove_duplicate=remove_duplicate)
        yield from gen

核心逻辑调用_named_members方法：

    def _named_members(self, get_members_fn, prefix='', recurse=True, remove_duplicate: bool = True):
        r"""Help yield various names + members of modules."""
        memo = set()
        modules = self.named_modules(prefix=prefix, remove_duplicate=remove_duplicate) if recurse else [(prefix, self)]
        for module_prefix, module in modules:
            members = get_members_fn(module)
            for k, v in members:
                if v is None or v in memo:
                    continue
                if remove_duplicate:
                    memo.add(v)
                name = module_prefix + ('.' if module_prefix else '') + k
                yield name, v

循环所有模块，获取每个模块的members，并返回。
get_members_fn其实就是调用_named_members时的第一个参数：module._parameters.items()，也就是module的所有参数。

总结

可以说，整个网络模型都是基于nn.Module类构建起来的，有两个核心方法：init和forward。Module可以嵌套组成一个复杂的网络结构， 无论是容器、还是模型、还是激活函数，都是一个Module的子类，这些Module环环相扣，共同完成一个可大可小的任务。
整体架构看似复杂，又可以很优雅。