PyTorch：模型定义

模型在深度学习中扮演着重要的角色，好的模型极大地促进了深度学习的发展进步，比如CNN的提出解决了图像、视频处理中的诸多问题，RNN/LSTM模型解决了序列数据处理的问题，GNN在图模型上发挥着重要的作用。当我们在向他人介绍一项深度学习工作的时候，对方可能首先要问的就是使用了哪些模型。

在本文中，我们将首先介绍三种基本的模型定义方式、然后介绍怎么使用模型快更快地搭建复杂网络（从零开始）。为了降低工作量和利用他人的研究成果，我们可以去修改现有的模型（可能非常大）来解决我们自己的问题。最后，我们介绍了模型的保存与读取，这在训练大规模模型时是非常有必要的。

1. 三种基本的模型定义方式

1.1 Sequential

当模型的前向计算为简单串联各个层的计算时， Sequential 只需要将模型的层按序排列起来即可。根据层名的不同，排列的时候有两种方式：

1.1.1 直接排列

import torch.nn as nn
net = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10), 
        )
print(net)

# Sequential(
#   (0): Linear(in_features=784, out_features=256, bias=True)
#   (1): ReLU()
#   (2): Linear(in_features=256, out_features=10, bias=True)
# )

1.1.2 使用 OrderedDict

import collections   # 实现了特定目标的容器，以提供Python标准内建容器 dict、list、set、tuple 的替代选择
import torch.nn as nn
net2 = nn.Sequential(collections.OrderedDict([   # 字典的子类，保留了他们被添加的顺序
          ('fc1', nn.Linear(784, 256)),
          ('relu1', nn.ReLU()),
          ('fc2', nn.Linear(256, 10))
          ]))
print(net2)

# Sequential(
#   (fc1): Linear(in_features=784, out_features=256, bias=True)
#   (relu1): ReLU()
#   (fc2): Linear(in_features=256, out_features=10, bias=True)
)

使用Sequential定义模型的好处在于简单、易读，同时使用Sequential定义的模型不需要再写forward，因为顺序已经定义好了。但使用Sequential也会使得模型定义丧失灵活性，比如需要在模型中间加入一个外部输入时就不适合用Sequential的方式实现。

1.2 ModuleList

ModuleList 接收一个子模块（或层，需属于nn.Module类）的列表作为输入，然后也可以类似List那样进行append和extend操作。

net = nn.ModuleList([nn.Linear(784, 256), nn.ReLU()])
net.append(nn.Linear(256, 10)) # 类似List的append操作
print(net[-1])  # 类似List的索引访问
print(net)

# Linear(in_features=256, out_features=10, bias=True)
# ModuleList(
#   (0): Linear(in_features=784, out_features=256, bias=True)
#   (1): ReLU()
#   (2): Linear(in_features=256, out_features=10, bias=True)
# )

要特别注意的是，nn.ModuleList 并没有定义一个网络，它只是将不同的模块储存在一起。ModuleList中元素的先后顺序并不代表其在网络中的真实位置顺序，需要经过forward函数指定各个层的先后顺序后才算完成了模型的定义。

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])

    def forward(self, x):
        # ModuleList can act as an iterable, or be indexed using ints
        for i, l in enumerate(self.linears):  # i 为索引（默认从0开始），l 为迭代的数据（nn.Linear）
            x = self.linears[i // 2](x) + l(x)   # i//2 指整数除法，可基于索引选择性调用不同位置的模块多次
        return x

1.3 ModuleDict

ModuleDict和ModuleList的作用类似，只是ModuleDict能够更方便地为神经网络的层添加名称。

net = nn.ModuleDict({
    'linear': nn.Linear(784, 256),
    'act': nn.ReLU(),
})
net['output'] = nn.Linear(256, 10) # 添加
print(net['linear']) # 访问
print(net.output)
print(net)

# Linear(in_features=784, out_features=256, bias=True)
# Linear(in_features=256, out_features=10, bias=True)
# ModuleDict(
#   (act): ReLU()
#   (linear): Linear(in_features=784, out_features=256, bias=True)
#   (output): Linear(in_features=256, out_features=10, bias=True)
# )

ModuleList和ModuleDict在某个完全相同的层需要重复出现多次时，非常方便实现，可以”一行顶多行“；

当我们需要之前层的信息的时候，比如 ResNets 中的残差计算，当前层的结果需要和之前层中的结果进行融合，一般使用 ModuleList/ModuleDict 比较方便。

2. 利用模型块快速搭建复杂网络

当模型的深度非常大时候，使用Sequential等定义模型结构需要向其中添加几百行代码，使用起来不甚方便。

对于大部分模型结构（比如ResNet、DenseNet等），虽然模型有很多层，但是其中有很多重复出现的结构。考虑到每一层有其输入和输出，若干层串联成的”模块“也有其输入和输出，如果我们能将这些重复出现的层定义为一个”模块“，每次只需要向网络中添加对应的模块来构建模型，这样将会极大便利模型构建的过程。

本节我们将以U-Net为例，介绍如何构建模型块，以及如何利用模型块快速搭建复杂模型。

2.1 U-Net简介

U-Net是分割 (Segmentation) 模型的杰作，在以医学影像为代表的诸多领域有着广泛的应用。U-Net模型结构如下图所示：

这个结构就是先对图片进行卷积和池化，然后进行通道的拼接concat，再对拼接之后的特征图做卷积和上采样，可以得到一个与输入图像尺寸相同的预测结果。

这个结构通过残差连接结构解决了模型学习中的退化问题（每一次下采样提炼特征的同时，也必然会损失一些边缘特征，而失去的特征并不能从上采样中找回，因此通过特征的拼接，来实现边缘特征的一个找回），使得神经网络的深度能够不断扩展。

图中每个模块上方的数目表示 channel_num，也就是 filter 的数量。Crop（裁剪）的目的是删除掉输入图像中无意义的像素，使其特征图大小与将要拼接的特征图保持一致。

Up sampling：由于输入图像通过卷积神经网络(CNN)提取特征后，输出的尺寸往往会变小，而有时我们需要将图像恢复到原来的尺寸以便进行进一步的计算(e.g.:图像的语义分割)，这个采用扩大图像尺寸，实现图像由小分辨率到大分辨率的映射的操作，叫做上采样。上采样有3种常见的方法：双线性插值(bilinear)，反卷积(Transposed Convolution)，反池化(Unpooling)。

2.2 U-Net模型块分析

结合上图，不难发现U-Net模型具有非常好的对称性。模型从上到下分为若干层，每层由左侧和右侧两个模型块组成，每侧的模型块与其上下模型块之间有连接；同时位于同一层左右两侧的模型块之间也有连接，称为“Skip-connection”。此外还有输入和输出处理等其他组成部分。由于模型的形状非常像英文字母的“U”，因此被命名为“U-Net”。

组成U-Net的模型块主要有如下几个部分：

• 每个子块内部的两次卷积（Double Convolution）
• 左侧模型块之间的下采样连接，即最大池化（Max pooling）
• 右侧模型块之间的上采样连接（Up sampling）
• 输出层的处理

除模型块外，还有模型块之间的横向连接，输入和U-Net底部的连接等计算，这些单独的操作可以通过forward函数来实现。

下面我们用PyTorch先实现上述的模型块，然后再利用定义好的模型块构建U-Net模型。

2.3 U-Net模型块实现

对应上一节分析的四个模型块，根据功能我们将其命名为：DoubleConv, Down, Up, OutConv。下面给出U-Net中模型块的 PyTorch 实现

import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),   # The mean and standard-deviation are calculated per-dimension over the mini-batches
            nn.ReLU(inplace=True),   # 当inplace=True的时候，会改变输入数据；当inplace=False的时候，不会改变输入数据
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)  # 因为 UNet 两次卷积使用的 out_channels 都相等，所以这里并未设置 mid_channels 参数；若自己的模型有不同 channel 的需求，这可以设置 mid_channels 参数。
        )

    def forward(self, x):
        return self.maxpool_conv(x)

class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)  # 待解决，为什么要加入一个 middle_channels 的输入参数；这可能与所使用的上采样方法的不同有关系
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)  # 反卷积，另外一种上采样方法
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):  # x2 指 skip-connection 输入的值，在图中的 size 要比 x1 大
        x1 = self.up(x1) # x1 上采样后，x1 与 x2 的长宽差距会较小，但 x2 还是会略大于 x1
        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]  # 计算 x1 和 x2 的长宽差距，便于通过填充等操作保持一致
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,  # 如果 diffX 是偶数，则左右填充的值是相等的，这样设置是为了处理 diffX 或 diffY 为奇数的情况
                        diffY // 2, diffY - diffY // 2])
        # The padding size by which to pad some dimensions of input are described starting from the last dimension and moving forward
        # https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html#torch.nn.functional.pad
      
        x = torch.cat([x2, x1], dim=1)   # 维度从 0 开始，维度 1 指的是 channel
        return self.conv(x)

class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__() 
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)   # 相当于全连接层的运算

    def forward(self, x):
        return self.conv(x)

2.4 利用模型块组装U-Net

使用写好的模型块，可以非常方便地组装 U-Net 模型。可以看到，通过模型块的方式实现了代码复用，整个模型结构定义所需的代码总行数明显减少，代码可读性也得到了提升。

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        factor = 2 if bilinear else 1
        self.down4 = Down(512, 1024 // factor)
        self.up1 = Up(1024, 512 // factor, bilinear)
        self.up2 = Up(512, 256 // factor, bilinear)
        self.up3 = Up(256, 128 // factor, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):  # example，按照 UNet 图的输入维度 x = torch.ones([10,1,572,572]); bilinear = True
        x1 = self.inc(x)   # torch.Size([10, 64, 572, 572])
        x2 = self.down1(x1)  # torch.Size([10, 128, 286, 286])
        x3 = self.down2(x2)  # torch.Size([10, 256, 143, 143])
        x4 = self.down3(x3)  # torch.Size([10, 512, 71, 71])
        x5 = self.down4(x4)  # torch.Size([10, 512, 35, 35])
        x = self.up1(x5, x4)  # torch.Size([10, 256, 71, 71])
        x = self.up2(x, x3)  # torch.Size([10, 128, 143, 143])
        x = self.up3(x, x2)  # torch.Size([10, 64, 286, 286])
        x = self.up4(x, x1)  # torch.Size([10, 64, 572, 572])
        logits = self.outc(x)  # torch.Size([10, 64, 572, 572])
        return logits

可以看到在上采样后，特征图的大小又恢复成了原来的大小，可以达到像素级别的预测。

注意：每一层得到的特征维度与 UNet 的图展示的不一样，这是因为实现代码并不是针对 UNet 的。UNet 的介绍可参考文章 https://blog.csdn.net/qq_39942341/article/details/119896754

3. 修改已有的模型

除了自己构建PyTorch模型外，还有另一种应用场景：我们已经有一个现成的模型，但该模型中的部分结构不符合我们的要求，为了使用模型，我们需要对模型结构进行必要的修改。

我们以pytorch官方视觉库torchvision预定义好的模型ResNet50为例来演示修改模型若干层、添加额外输入和额外输出。

3.1 修改模型层

操作步骤：载入已有的模型 → 定义新的层 → 赋值/替换（类似于替换张量或列表中的某个值）

import torchvision.models as models
net = models.resnet50()
print(net)

# ResNet(
#   (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
#   (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
#   (relu): ReLU(inplace=True)
#   (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
#   (layer1): Sequential(
#     (0): Bottleneck(
#       (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
#       (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
#       (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
#       (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
#       (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
#       (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
#       (relu): ReLU(inplace=True)
#       (downsample): Sequential(
#         (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
#         (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
#       )
#     )
..............
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=2048, out_features=1000, bias=True)
)

现有模型的全连接层（fc）的输出节点数是1000。

假设我们要用这个resnet模型去做一个10分类的问题，就应该修改模型的fc层，将其输出节点数替换为10。另外，我们觉得一层全连接层可能太少了，想再加一层。可以做如下修改：

from collections import OrderedDict
classifier = nn.Sequential(OrderedDict([('fc1', nn.Linear(2048, 128)),
                          ('relu1', nn.ReLU()), 
                          ('dropout1',nn.Dropout(0.5)),
                          ('fc2', nn.Linear(128, 10)),
                          ('output', nn.Softmax(dim=1))
                          ]))
    
net.fc = classifier

这里的操作相当于将模型（net）最后名称为“fc”的层替换成了名称为“classifier”的结构，该结构是我们自己定义的。

3.2 添加外部输入

操作步骤：定义新的输入 → 将新的输入与目标层的已有输入合并（可能是维度或者是值处理）→ 合并输入作为目标层的最终输入

我们希望利用已有的模型结构，在倒数第二层增加一个额外的输入变量add_variable来辅助预测。具体实现如下：

class Model(nn.Module):
    def __init__(self, net):
        super(Model, self).__init__()
        self.net = net
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc_add = nn.Linear(1001, 10, bias=True)  # 维度处理，1000维增加了一维，也就是数据拼接处理
        self.output = nn.Softmax(dim=1)
        
    def forward(self, x, add_variable):
        x = self.net(x)
        x = torch.cat((self.dropout(self.relu(x)), add_variable.unsqueeze(1)),1)  # add_variable 和 原本的输入 拼接在一起
        x = self.fc_add(x)
        x = self.output(x)
        return x

3.3 添加额外输出

操作步骤：将要输出的变量用不同命名 → forward 函数中 return 多个变量

在已经定义好的模型结构上，同时输出1000维的倒数第二层和10维的最后一层结果。具体实现如下：

class Model(nn.Module):
    def __init__(self, net):
        super(Model, self).__init__()
        self.net = net
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc1 = nn.Linear(1000, 10, bias=True)
        self.output = nn.Softmax(dim=1)
        
    def forward(self, x, add_variable):
        x1000 = self.net(x)
        x10 = self.dropout(self.relu(x1000))
        x10 = self.fc1(x10)
        x10 = self.output(x10)
        return x10, x1000

4. 模型保存和读取

一个PyTorch模型主要包含两个部分：模型结构和权重。其中模型是继承nn.Module的类，权重的数据结构是一个字典（key是层名，value是权重向量）。存储也由此分为两种形式：存储整个模型（包括结构和权重），和只存储模型权重。

from torchvision import models
model = models.resnet152(pretrained=True)

# 保存整个模型
torch.save(model, save_dir)
# 保存模型权重
torch.save(model.state_dict, save_dir)

对于PyTorch而言，pt, pth和pkl三种数据格式均支持模型权重和整个模型的存储，因此使用上没有差别。

因为作者没有 GPU，故目前只记录使用 CPU 的读取方式：

# 读取整个模型
model = torch.load(save_path, map_location='cpu')

# 读取模型权重
loaded_dict = torch.load(save_dir)
model.state_dict = loaded_dict  # model 需要事先定义好，与模型权重对应的结构一致

参考资料

• 图像分割必备知识点 | Unet详解 理论+ 代码
• 三种上采样方法 | Three up sampling methods
• DataWhale教程：Pytorch模型定义
• UNet学习