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 需要事先定义好，与模型权重对应的结构一致
