PyTorch05-PyTorch模型定义
五、PyTorch模型定义
5.1 PyTorch模型定义的方式
我们会在不同场合看到多种不同的模型架构，比如 CNN（Convolutional Neural Network，卷积神经网络）、RNN（Recurrent Neural Network，循环神经网络）、GNN（Graph Neural Network，图神经网络），它们是深度学习领域针对不同数据结构设计的神经网络架构，分别用来解决网格化结构数据（如二维图像、视频帧）的局部特征提取与空间感知问题、序列数据（如文本、语音、时间序列）的前后关联与时序依赖问题、以及非欧几里得图结构数据（如社交网络、分子结构、知识图谱）的复杂关系建模问题。
在实际应用中，每种宏观架构下都有诸多经典的具体模型实现。在学习或使用这些模型时，通常需要先理解其内部的运行机制。本章节将基于 PyTorch 这一深度学习框架，从粗略地角度来探索与学习下模型的代码实现与构造方式。
本章节学习的内容包括：使用PyTorch定义模型的三种经典方式。
5.1.1 必要的知识回顾
在PyTorch
中，torch.nn.Module是一个基本构造类，所有神经网络层（如 nn.Linear、nn.Conv2d）以及自定义的大模型，本质上都是它的子类。
当我们基于torch.nn.Module来自定义模型时，一般就是要重写自定义类的__init__与forward两个方法。
在基于torch.nn.Module来构建自定义模型组成与执行流程的具体实现时，可以借助Sequential或
ModuleList或ModuleDict
这三个类（它们是官方提供且最常用的，但并不是 PyTorch
定义模型子模块的“唯三”方式）。其他的常用的方式还有：给自定义类的属性赋值然后在forward方法中调用。通常在编写较复杂的网络时，是将“直接类属性定义”与
Sequential、ModuleList、ModuleDict
混合配合使用的。
以下分别对这三个类及使用进行精简介绍并给出示例。 ### 5.1.2 定义模型子模块的常用三个类
5.1.1.2 Sequential
nn.Sequential 是 PyTorch
中最简单的模型容器。它按照层添加的顺序依次执行前向传播，无需手动编写
forward()
函数。因此它无法应对条件分支、跳连、多分支等复杂结构，只适合于结构简单、顺序执行的模型。
它可以接收一系列子模块或子模块的有序字典(OrderedDict)，从而给模型逐一添加
Module 的实例，且按它们出现的顺序作为模型前向传播时的执行顺序。
直接添加子模块与子模块的有序字典(OrderedDict)的用法大致如下：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15# 方式一：直接传入多个 nn.Module（匿名层，无名称）
net1 = nn.Sequential(
nn.Linear(20, 128), # 第 0 层：输入 20 → 输出 128
nn.ReLU(), # 第 1 层：ReLU 激活
nn.Linear(128, 10), # 第 2 层：输入 128 → 输出 10
nn.Softmax(dim=1), # 第 3 层：Softmax 概率分布
)
# 方式二：使用 OrderedDict 给每层命名（推荐）
net2 = nn.Sequential(collections.OrderedDict([
('fc1', nn.Linear(20, 128)), # 命名：fc1（第一个全连接层）
('relu1', nn.ReLU()), # 命名：relu1
('fc2', nn.Linear(128, 10)), # 命名：fc2
('softmax', nn.Softmax(dim=1)), # 命名：softmax
]))
完整可执行的示例如下：
1 | |
使用Sequential虽然方便，但也会使得模型定义丧失灵活性，如需在模型中间加入一个外部输入，就不适合用Sequential的方式实现。
5.1.1.3 ModuleList
nn.ModuleList 接收一个子模块（或层，要是nn.Module类）的列表作为输入，但nn.ModuleList并没有直接定义一个完整的模型（无法自动进行前向传播），它只是将不同的子模块或层储存在一起，并将子模块或层的权重也会自动添加到模型中。要想让它像模型一样工作还需手动遍历其中所有层或子模块（即列表中所有元素）并执行它们。
1 | |
对于一个nn.ModuleList 实例，可以操作Python中普通的列表List一样对其增加、扩展、删除、索引与修改元素。
1
2
3# append：像普通 list 一样往 ModuleList 末尾追加一个层
net.append(nn.Softmax(dim=1)) # 在末尾追加 Softmax，将输出转为概率分布
print(f"追加 Softmax 后，ModuleList 共包含 {len(net)} 个层")
完整的可执行示例如下： 1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196# -*- coding: utf-8 -*-
"""
02-使用ModuleList定义模型.py
============================
ModuleList 是 PyTorch 提供的一个「列表式」模块容器。
它的作用类似于 Python 的 list，专门用来存放多个 nn.Module 子层。
与普通 list 的核心区别：
- 普通 list 存 nn.Module 不会被 PyTorch 自动识别为模型参数，
导致 model.parameters() 无法正确收集这些层的权重。
- ModuleList 会正确注册所有子模块，参数可以被自动收集。
本示例：
1. 直接使用 ModuleList 构建一个简单的两层全连接网络
2. 用随机数据跑一次前向传播，观察输出形状
3. 对比 ModuleList 和普通 list 在参数收集上的区别
"""
import torch
import torch.nn as nn
# ============================================================
# 第一部分：直接使用 ModuleList（不定义类）
# ============================================================
print("=" * 50)
print("第一部分：直接使用 ModuleList")
print("=" * 50)
# 用 ModuleList 定义两个全连接层 + 一个激活函数
# 注意：这里只存了层，还没有指定前向传播的顺序
net = nn.ModuleList([
nn.Linear(20, 128), # 第一层：输入 20 维 → 输出 128 维
nn.ReLU(), # 激活函数：逐元素应用 ReLU
nn.Linear(128, 10), # 第二层：输入 128 维 → 输出 10 维（10 个类别）
])
# append：像普通 list 一样往 ModuleList 末尾追加一个层
net.append(nn.Softmax(dim=1)) # 在末尾追加 Softmax，将输出转为概率分布
print(f"追加 Softmax 后，ModuleList 共包含 {len(net)} 个层")
# 索引访问：像普通 list 一样通过下标取层
last_layer = net[-1]
print(f"最后一层是：{last_layer}")
# ModuleList 本身只负责管理子模块，不负责前向传播的顺序。
# 前向传播时，需要手动遍历 ModuleList 中的每一层。
print("\n--- 前向传播 ---")
# 构造一个 batch 的随机输入：batch_size=4，每个样本 20 维
dummy_input = torch.randn(4, 20)
print(f"输入形状：{dummy_input.shape}")
params = sum(p.numel() for p in net.parameters())
print(f" net 的可训练参数数量：{params} ")
# 手动遍历每一层，依次执行前向计算
x = dummy_input
for i, layer in enumerate(net):
x = layer(x)
print(f" 经过第 {i} 层 {type(layer).__name__:12s} 后，形状变为：{tuple(x.shape)}")
print(f"最终输出形状：{x.shape}") # 期望：(4, 10)
# ============================================================
# 第二部分：在 nn.Module 子类中使用 ModuleList
# 这是 ModuleList 最典型的用法
# ============================================================
print("\n" + "=" * 50)
print("第二部分：在 nn.Module 中使用 ModuleList")
print("=" * 50)
class MyModel(nn.Module):
"""
一个使用 ModuleList 的两层全连接分类模型。
网络结构：
输入(20维) → Linear(20→128) → ReLU → Linear(128→10) → Softmax
"""
def __init__(self):
super().__init__() # 必须调用父类 nn.Module 的构造函数
# 将所有层放入 ModuleList
self.layers = nn.ModuleList([
nn.Linear(20, 128), # 隐藏层：20 → 128
nn.ReLU(), # 激活
nn.Linear(128, 10), # 输出层：128 → 10
nn.Softmax(dim=1), # 输出概率分布
])
def forward(self, x):
"""
前向传播：手动遍历 self.layers，逐层执行。
"""
for layer in self.layers:
x = layer(x)
return x
model = MyModel()
print(f"模型结构：\n{model}")
# 用随机数据跑一次前向传播
sample = torch.randn(4, 20)
output = model(sample)
print(f"\n输入形状：{tuple(sample.shape)}")
print(f"输出形状：{tuple(output.shape)}") # 期望：(4, 10)
print(f"输出示例（第一个样本）：{output[0].detach().cpu().numpy().round(4)}")
print(f"概率之和（应≈1.0）：{output.sum(dim=1).detach().cpu().numpy().round(4)}")
# ============================================================
# 第三部分：ModuleList vs 普通 list（参数收集对比）
# ============================================================
print("\n" + "=" * 50)
print("第三部分：ModuleList vs 普通 list — 参数收集对比")
print("=" * 50)
class ModelWithPlainList(nn.Module):
"""使用普通 Python list 存放子模块（不推荐）"""
def __init__(self):
super().__init__()
# 错误示范：普通 list 不会被 PyTorch 正确注册
self.wrong_list = [
nn.Linear(20, 128),
nn.ReLU(),
nn.Linear(128, 10),
]
def forward(self, x):
for layer in self.wrong_list:
x = layer(x)
return x
class ModelWithModuleList(nn.Module):
"""使用 ModuleList 存放子模块（推荐）"""
def __init__(self):
super().__init__()
self.right_list = nn.ModuleList([
nn.Linear(20, 128),
nn.ReLU(),
nn.Linear(128, 10),
])
def forward(self, x):
for layer in self.right_list:
x = layer(x)
return x
wrong = ModelWithPlainList()
right = ModelWithModuleList()
# model.parameters() 会递归收集所有可训练参数（权重 + 偏置）
wrong_params = sum(p.numel() for p in wrong.parameters())
right_params = sum(p.numel() for p in right.parameters())
print(f"普通 list 模型的可训练参数数量：{wrong_params} （应为 0，层没被注册）")
print(f"ModuleList 模型的可训练参数数量：{right_params} （两个 Linear 层权重+偏置）")
# ============================================================
# 第四部分：ModuleList 常见操作速查
# ============================================================
print("\n" + "=" * 50)
print("第四部分：ModuleList 常见操作速查")
print("=" * 50)
mylist = nn.ModuleList([
nn.Linear(10, 20),
nn.ReLU(),
])
# 追加单个层
mylist.append(nn.Linear(20, 5))
print(f"append 后长度：{len(mylist)}")
# 扩展多个层
mylist.extend([nn.ReLU(), nn.Softmax(dim=1)])
print(f"extend 后长度：{len(mylist)}")
# 弹出最后一个（ModuleList.pop() 需要传入索引，不像普通 list 无参弹出）
popped = mylist.pop(-1)
print(f"pop 弹出：{type(popped).__name__}，剩余长度：{len(mylist)}")
# 插入到指定位置
mylist.insert(1, nn.Dropout(0.5))
print(f"insert 后第 1 层：{type(mylist[1]).__name__}")
# 删除指定位置
del mylist[1]
print(f"del 后第 1 层：{type(mylist[1]).__name__}")
print("\nModuleList 示例执行完毕！")
5.1.1.4 ModuleDict
nn.ModuleDict和nn.ModuleList的作用类似，只是nn.ModuleDict的结构是一个加强版的Dict，能够给神经网络的模块或层设置名称（其实就是Dict中Key），也支持一系列Dict中常见操作如添加、更新、删除等。
1
2
3
4
5
6
7
8
9
10
11
12
13# 创建 ModuleDict，用 dict 初始化（方式一：传入字典）
net = nn.ModuleDict({
'linear': nn.Linear(20, 128), # 键 'linear' 对应一个全连接层
'act': nn.ReLU(), # 键 'act' 对应一个激活层
})
# 动态添加新的层（方式二：像普通 dict 一样用 [] 赋值）
net['output'] = nn.Linear(128, 10) # 添加键 'output'
# 手动指定前向顺序：linear → act → output
x = net['linear'](x)
x = net['act'](x)
x = net['output'](x)1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267# -*- coding: utf-8 -*-
"""
03-使用ModuleDict定义模型.py
=============================
ModuleDict 是 PyTorch 提供的一个「字典式」模块容器。
它的作用类似于 Python 的 dict，专门用来存放多个 nn.Module 子层，并允许通过名称（字符串）来访问。
与普通 dict 的核心区别：
- 普通 dict 存 nn.Module 不会被 PyTorch 自动识别为模型参数，
导致 model.parameters() 无法正确收集这些层的权重。
- ModuleDict 会正确注册所有子模块，参数可以被自动收集。
ModuleDict 的典型使用场景：
- 需要动态选择某几层（例如：根据不同条件选用不同的卷积层）
- 需要像查字典一样通过名字获取某层（例如 net['backbone']）
- 需要保存/加载时通过名称管理模块（例如：保存为 .pt 文件时方便按名称恢复）
本示例：
1. 构建一个简单的两层网络
2. 演示 ModuleDict 的基本操作（增/查/删/遍历）
3. 用随机数据跑前向传播
4. 对比 ModuleDict 和普通 dict 在参数收集上的区别
"""
import torch
import torch.nn as nn
# ============================================================
# 第一部分：ModuleDict 基础用法
# ============================================================
print("=" * 50)
print("第一部分：ModuleDict 基础用法")
print("=" * 50)
# 创建 ModuleDict，用 dict 初始化（方式一：传入字典）
net = nn.ModuleDict({
'linear': nn.Linear(20, 128), # 键 'linear' 对应一个全连接层
'act': nn.ReLU(), # 键 'act' 对应一个激活层
})
# 动态添加新的层（方式二：像普通 dict 一样用 [] 赋值）
net['output'] = nn.Linear(128, 10) # 添加键 'output'
print(f"ModuleDict 共包含 {len(net)} 个模块")
print(f"各模块的键：{list(net.keys())}")
print(f"各模块的值：{list(net.values())}")
# 通过名称访问某层
print(f"\n通过 net['linear'] 访问：{net['linear']}")
print(f"通过 net.output 访问：{net.output}")
# 前向传播：ModuleDict 不自动按顺序执行，需要在 forward() 中手动调用
print("\n--- 前向传播 ---")
dummy_input = torch.randn(4, 20)
print(f"输入形状：{dummy_input.shape}")
# 手动指定前向顺序：linear → act → output
x = net['linear'](dummy_input)
x = net['act'](x)
x = net['output'](x)
print(f"最终输出形状：{x.shape}") # 期望：(4, 10)
# ============================================================
# 第二部分：在 nn.Module 子类中使用 ModuleDict（推荐方式）
# 这是 ModuleDict 最典型的用法
# ============================================================
print("\n" + "=" * 50)
print("第二部分：在 nn.Module 中使用 ModuleDict")
print("=" * 50)
class MyModel(nn.Module):
"""
一个使用 ModuleDict 的两层全连接分类模型。
网络结构：
输入(20维) → linear(20→128) → ReLU → output(128→10) → Softmax
"""
def __init__(self):
super().__init__()
# 将各层放入 ModuleDict，通过名称管理
self.layers = nn.ModuleDict({
'hidden': nn.Linear(20, 128), # 隐藏层
'act': nn.ReLU(), # 激活层
'out': nn.Linear(128, 10), # 输出层
})
def forward(self, x):
# 通过名称逐层调用
x = self.layers['hidden'](x)
x = self.layers['act'](x)
x = self.layers['out'](x)
return x
model = MyModel()
print(f"模型结构：\n{model}")
# 前向传播
sample = torch.randn(4, 20)
output = model(sample)
print(f"\n输入形状：{tuple(sample.shape)}")
print(f"输出形状：{tuple(output.shape)}")
# ============================================================
# 第三部分：ModuleDict 常见操作速查
# ============================================================
print("\n" + "=" * 50)
print("第三部分：ModuleDict 常见操作速查")
print("=" * 50)
md = nn.ModuleDict()
md['a'] = nn.Linear(10, 20) # 添加：md['key'] = module
print(f"添加 'a' 后，长度 = {len(md)}")
md.update({'b': nn.ReLU()}) # 此处更新其实为添加
print(f"update {{'b':...}} 后，长度 = {len(md)}")
print(f"md['a']：{md['a']}") # 访问：md['key']
print(f"md.a：{md.a}") # 也可用属性方式访问
# 遍历所有键值对
print("\n遍历所有模块：")
for name, module in md.items():
print(f" {name}: {module}")
# 删除模块
del md['a']
print(f"\n删除 'a' 后，键列表 = {list(md.keys())}")
# pop 删除并返回模块
popped = md.pop('b')
print(f"pop 弹出 'b'：{type(popped).__name__}")
print(f"pop 后长度 = {len(md)}")
# clear 清空所有模块
md.clear()
print(f"clear 后长度 = {len(md)}")
# ============================================================
# 第四部分：ModuleDict vs 普通 dict — 参数收集对比
# ============================================================
print("\n" + "=" * 50)
print("第四部分：ModuleDict vs 普通 dict — 参数收集对比")
print("=" * 50)
class ModelWithPlainDict(nn.Module):
"""使用普通 Python dict 存放子模块（不推荐）"""
def __init__(self):
super().__init__()
# 错误示范：普通 dict 不会被 PyTorch 正确注册
self.wrong_dict = {
'linear': nn.Linear(20, 128),
'act': nn.ReLU(),
'out': nn.Linear(128, 10),
}
def forward(self, x):
x = self.wrong_dict['linear'](x)
x = self.wrong_dict['act'](x)
x = self.wrong_dict['out'](x)
return x
class ModelWithModuleDict(nn.Module):
"""使用 ModuleDict 存放子模块（推荐）"""
def __init__(self):
super().__init__()
self.right_dict = nn.ModuleDict({
'linear': nn.Linear(20, 128),
'act': nn.ReLU(),
'out': nn.Linear(128, 10),
})
def forward(self, x):
x = self.right_dict['linear'](x)
x = self.right_dict['act'](x)
x = self.right_dict['out'](x)
return x
wrong = ModelWithPlainDict()
right = ModelWithModuleDict()
wrong_params = sum(p.numel() for p in wrong.parameters())
right_params = sum(p.numel() for p in right.parameters())
print(f"普通 dict 模型的可训练参数数量：{wrong_params} （应为 0，层没被注册）")
print(f"ModuleDict 模型的可训练参数数量：{right_params} （两个 Linear 层权重+偏置）")
# ============================================================
# 第五部分：ModuleDict 实战示例——按名称选择不同层
# 这是 ModuleDict 最强大的用法：
# 可以用字符串动态决定使用哪一层
# ============================================================
print("\n" + "=" * 50)
print("第五部分：ModuleDict 实战——动态选择分支")
print("=" * 50)
class MultiBranchModel(nn.Module):
"""
一个具有多分支结构的模型，使用 ModuleDict 管理分支。
场景：输入有两种处理方式（简单路径 / 复杂路径），
根据输入特征动态选择走哪条路径。
"""
def __init__(self):
super().__init__()
# 两个分支用 ModuleDict 管理
self.branches = nn.ModuleDict({
'simple': nn.Sequential( # 简单分支：单层
nn.Linear(20, 10),
),
'complex': nn.Sequential( # 复杂分支：两层 + 激活
nn.Linear(20, 64),
nn.ReLU(),
nn.Linear(64, 10),
),
})
def forward(self, x, branch_name='complex'):
# 根据传入的 branch_name 字符串，动态选择使用哪个分支
if branch_name not in self.branches:
raise ValueError(f"未知分支 '{branch_name}'，可选：{list(self.branches.keys())}")
return self.branches[branch_name](x)
mb = MultiBranchModel()
print(f"可选分支：{list(mb.branches.keys())}")
sample = torch.randn(4, 20)
# 走简单分支
out_simple = mb(sample, branch_name='simple')
print(f"\n简单分支输出形状：{tuple(out_simple.shape)}")
# 走复杂分支
out_complex = mb(sample, branch_name='complex')
print(f"复杂分支输出形状：{tuple(out_complex.shape)}")
# ============================================================
# 总结
# ============================================================
print("\n" + "=" * 50)
print("ModuleDict 使用要点总结")
print("=" * 50)
summary = """
1. nn.ModuleDict() 创建一个字典式的模块容器。
2. 通过字符串键名访问模块：md['name'] 或 md.name。
3. 支持字典操作：md['key'] = module、update()、pop()、clear()、del 等。
4. 适合需要「按名称动态选择/管理模块」的场景。
5. ModuleDict 本身不自动执行前向传播，需在 forward() 中手动调用各层。
6. 适用场景：
- 多分支网络（根据条件选不同分支）
- 需要保存/加载时按名称恢复模块
- 可插拔的层结构（如更换 backbone）
"""
print(summary)
print("ModuleDict 示例执行完毕！")
5.1.3 三种方法的比较与适用场景
nn.Sequential
适用于结构简单、单线顺序执行的网络或子网络。它的优势在于能自动处理层与层之间的数据传递，无需在
forward
中逐层手写连接逻辑；ß但无法直接支持跳连、多分支和动态条件判断。
nn.ModuleList 和 nn.ModuleDict
则适用于需要动态选择、遍历堆叠或分支切换的复杂网络结构。它们主要起
“安全注册参数”和“容器管理”
的作用，具体的跳连、分支以及数据流向逻辑，必须在自定义模型的
forward() 中手动编写实现。
| 容器类型 | 自动执行 forward | 适用场景 | 关键局限性 |
|---|---|---|---|
nn.Sequential |
是 | 单线流水线结构、无分支/无跳连、快速搭建局部子网络（如 MLP、经典 CNN 骨干） | 无法支持多输入/多输出、条件分支、跳连（如 ResNet） |
nn.ModuleList |
否 | 需要按索引访问层、循环堆叠大量相同 Block（如 Transformer 编码器层）、动态选择执行某些层 | 必须在自定义 forward 中用 for
循环手动调用每一层 |
nn.ModuleDict |
否 | 需要根据名称（字符串 Key）动态调用不同分支（如多任务学习中的不同 Head、多数据集适配器） | 必须在自定义 forward 中通过 Key
手动查找并调用对应模块 |
5.2 利用模型块快速搭建复杂网络
到目前为止，所学的知识下，都是利用torch.nn包下的相关类来直接定义模型，如Sequential、ModuleList、ModuleDict等。但当模型结构比较复杂、导数很多时，直接使用它们当然可以完成模型的定义，但层层不断地堆叠子模块或层会使得模型的可读性与复用程度很低。通过阅读前面学习过的多个规模的ResNet模型的具体实现可以发现，其中一直在复用两个模块：BasicBlock、Bottleneck，这就是模型块。
复杂、多层的模型中往往会重复出现类似的层或子模块，它们都会有输入与输出。将输入输出、结构重复或类似的层或子模块进行抽象，就形成了模型块。复杂、多层的模型一般利用模型块来进行实现，减少了代码量同时也提升了模型的可读性。
本章节以U-Net模型为例，讲述模型结构的特征或规律、如何构建模型块、如何利用模型块构建复杂模型（这些也就是本章节的学习内容）。
5.2.1 U-Net简介
U-Net是分割 (Segmentation) 模型的杰作，在以医学影像为代表的诸多领域有着广泛的应用。 笔者将U-Net这个开源项目的github代码仓库fork了一份：Pytorch-UNet，从此页面也可以找到项目官方链接。U-Net模型结构（当前master分支代码对应的模型结构）如下图所示。
5.2.2 U-Net模型结构分析
从上图可以看到，UNet模型的结构在垂直方向可以分为几层；除了最底层的其他的层，每层都有左右两个部分；对于中间层，左右部分都与其上层或下层的部分存在连接，对于最上层或最底层，仍在反方向与其他部分存在连接；对于中间层，左右两个部分之间还存在一条跳跃连接路径。UNet整个模型结构图像是一个U型，这也是此模型被称为UNet的缘由。
再明确大致分析下其结构特点： - 每层的左侧部分都会进行两次卷积操作。===>可封装成一个模型块DoubleConv。 - 每层中（除了第一层），与相邻的部分之间会进行一次最大池化以进行一次下采样，再连续进行两次卷积（可复用DoubleConv）。===>可封装成一个模型块Down。 - 每层中的右侧部分，都会进行一次上采样，并与同一层的左侧部分输出执行跳跃连接操作，最后再连续进行两次卷积操作。===>可封装成一个模型块Up。 - 模型的最后输出是用与多个(1,1)卷积核做卷积得到的每个分类的原始得分。===>可单独封装成一个子模型块OutConv。
这是对UNet模型的结构分析，及尝试封装成对应的模型块。其他实现细节此处不展开。
5.2.3 U-Net模型块实现
5.2.3.1 U-Net模型块实现代码
以下是对上述分析得到的4个模型块的实现（来自PyTorch-UNet开源项目的master分支），其中Class
Up中bilinear相关代码暂时不用关注（笔者暂时也去研究它的用处。执行train.py进行测试时，使用默认参数时没有用到）。
以下代码来自unet/unet_parts.py的实现，笔者给其添加上详细了注释。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127""" Parts of the U-Net model """
# 本文件定义 U-Net 的可复用基础模块：编码器中的双卷积与下采样、
# 解码器中的上采样与跳跃连接，以及将特征映射投影为最终类别通道的输出层。
import torch
import torch.nn as nn
import torch.nn.functional as F
# 两个连续的 3×3 卷积块。卷积均使用 padding=1，因此不会改变特征图的高和宽。
# 若输入为 [N, in_channels, H, W]，输出为 [N, out_channels, H, W]。
# 每次卷积后接批归一化和原地 ReLU；卷积层关闭 bias，因为 BatchNorm 已提供可学习偏置。
class DoubleConv(nn.Module):
"""(convolution => [BN] => ReLU) * 2"""
def __init__(self, in_channels, out_channels, mid_channels=None):
super().__init__()
# 默认让两次卷积之间的通道数等于输出通道数；上采样的双线性分支会显式传入
# mid_channels，以便先压缩拼接后的通道数并减少参数量。
if not mid_channels:
mid_channels = out_channels
self.double_conv = nn.Sequential(
# 第一次卷积：只改变通道数 in_channels → mid_channels，空间尺寸保持 H × W。
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(mid_channels),
nn.ReLU(inplace=True),
# 第二次卷积：将中间通道数变为目标通道数 out_channels。
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):
# 输入/输出的空间尺寸相同，通道数由 in_channels 变为 out_channels。
return self.double_conv(x)
# 编码器的一个下采样阶段：先将高、宽各减半，再提取更高层的局部特征。
# 输入 [N, in_channels, H, W]，输出 [N, out_channels, floor(H/2), floor(W/2)]。
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
# 默认 stride 等于 kernel_size=2，因此特征图尺寸缩小为原来的 1/2。
# 当 H 或 W 为奇数时，MaxPool2d 会向下取整；例如宽 959 会变成 479。
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
# 通道数通常在该阶段增加，以补偿空间分辨率下降带来的信息损失。
return self.maxpool_conv(x)
# 解码器的一个上采样阶段。x1 来自更深层的解码特征，x2 来自编码器的同尺度跳跃连接。
# 该模块完成四件事：上采样 x1、与 x2 对齐、沿通道维拼接、用 DoubleConv 融合两类特征。
# 对默认的转置卷积分支，输入满足：
# x1: [N, in_channels, H, W]
# x2: [N, in_channels // 2, 约 2H, 约 2W]
# 最终输出为 [N, out_channels, x2.size(2), x2.size(3)]。
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:
# 双线性插值将 x1 的高、宽各放大 2 倍，但不改变其通道数，也不引入可学习参数。
# 此时 x1 与 x2 拼接后共有 in_channels 个通道；DoubleConv 的第一层将它们
# 压缩到 in_channels // 2，第二层再输出 out_channels，从而控制参数量。
# align_corners=True 使输入和输出的角点坐标对齐，会影响插值的具体数值。
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
else: # 在train.py 文件中，bilinear参数使用默认的False值
# 默认分支。kernel_size=stride=2 且没有 padding 时，转置卷积会使 H、W 各变为 2 倍，
# 并把通道数 in_channels 减为 in_channels // 2；该变换的参数可通过训练学习。
# 拼接后的通道数重新变为 in_channels，因此后续 DoubleConv 可直接使用默认中间通道数。
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):
# x1 是较深层、分辨率较低但语义更强的特征；x2 是编码器保留下来的细节特征。
# 例如默认模型的 up1 中：x1 为 [N, 1024, H/16, W/16]，
# 上采样后为 [N, 512, H/8, W/8]，x2（down3 输出）也有 512 个通道。
x1 = self.up(x1) ## 针对up1: (1, 1024, 40, 59) -> (1, 512, 80, 118)
# input is CHW
# 实际张量采用 PyTorch 的 NCHW 布局；第 2、3 维分别是 H、W，故使用 size()[2:4]。
# 记录两个空间维度的差值。输入尺寸不是 2 的整数次幂时，编码和解码路径
# 的尺寸可能相差一个像素，需要在拼接前对齐。
# 例如当前 Carvana 默认输入宽 959：池化后为 479、239、119、59，
# 而 59 经 2× 转置卷积会得到 118；因此必须补 1 列才能与宽 119 的 x2 拼接。
diffY = x2.size()[2] - x1.size()[2] # 针对up1: diffY = 80 - 80 = 0
diffX = x2.size()[3] - x1.size()[3] # 针对up1: diffX = 119 - 118 = 1
# F.pad 对四维 NCHW 张量的填充顺序为 [左, 右, 上, 下]。将奇数差值
# 分到两侧：例如 diffX=1 时变为 [0, 1]，额外像素放在右侧；diffX=2 时为 [1, 1]。
# 经过填充后，x1 的空间尺寸与 x2 完全相同。PyTorch 对负填充值会裁剪，
# 但本网络的常规调用期望 x2 尺寸不小于上采样后的 x1。
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
# if you have padding issues, see
# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
# dim=1 是通道维而非空间维。拼接不混合像素，而是在每个空间位置将两个特征向量首尾相接：
# x2 [N, in_channels//2, H', W'] + x1 [N, in_channels//2, H', W']
# → x [N, in_channels, H', W']。
# 随后的 DoubleConv 才会学习如何融合编码器的定位细节与解码器的高层语义。
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
# 将最后一层特征映射逐像素投影到目标输出通道（例如语义分割的类别数）。
# 1×1 卷积只混合通道信息，不改变批大小和空间分辨率。
# 输入 [N, in_channels, H, W]，输出 [N, out_channels, H, W]；out_channels 即类别数。
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):
# 此处返回原始 logits；损失函数或推理代码负责按任务选择 sigmoid/softmax。
return self.conv(x)
5.2.3.2 转置卷积操作
这个实现中涉及到一个新的概念，叫转置卷积操作，也叫反卷积，也常被称作上采样。与前面经常使用的下采样基本上可以看作是互逆的操作（但非严格互逆）。
nn.ConvTranspose2d 是 PyTorch
中的二维转置卷积操作（Transposed
Convolution），在学术界和工程中也常被称为
“反卷积”（Deconvolution） 或
“上采样卷积”。
它的核心作用与标准卷积（nn.Conv2d）刚好相反：标准卷积通常用于“下采样”（提取特征、缩小图像尺寸），而转置卷积用于“上采样”（放大特征图分辨率、恢复空间细节）。
5.2.3.2.1 为什么叫“转置”卷积？
标准卷积在数学运算时，可以写成矩阵乘法 \(Y = C \cdot X\)（其中 \(C\) 是根据卷积核构建的稀疏矩阵）。 而转置卷积在数学上的前向传播，相当于使用了标准卷积矩阵的转置矩阵 \(C^T\) 进行计算（即 \(Y = C^T \cdot X\)），因此得名“转置卷积”。
5.2.3.2.2 工作原理（直观理解）
转置卷积并不是简单地把卷积过程逆转恢复出原始像素，而是一种带可学习参数的图像放大技术。
- 补白（Padding / Stride 扩展）：在输入特征图的元素之间插入空白（零填充），拉开像素间距。
- 卷积滑动：用可学习的卷积核在扩展后的特征图上进行标准卷积滑动。
- 输出放大：最终生成一个长宽尺寸更大的特征图。
5.2.3.2.3 维度计算公式
给定输入尺寸 \(H_{in}\)，转置卷积输出的尺寸 \(H_{out}\) 计算公式如下： \[H_{out} = (H_{in} - 1) \times \text{stride} - 2 \times \text{padding} + \text{kernel\_size} + \text{output\_padding}\] 最常见的 2 倍上采样配置（放大 2 倍）：
当设置 kernel_size=2, stride=2, padding=0 或
kernel_size=4, stride=2, padding=1
时，输出图像的长宽正好是输入的 2 倍。
1 | |
5.2.3.2.4 适用场景与局限性
- 经典应用场景
- UNet / FPN：在 Segmentation 的 Decoder 部分用于逐步恢复图像分辨率。
- GAN（生成对抗网络）：生成器（Generator）用来从低维随机噪声逐步生成高分辨率图像。
- Diffusion（扩散模型）：在主干 UNet 的上采样阶段放大特征图。
- 局限性：棋盘效应（Checkerboard Artifacts）
- 由于转置卷积的卷积核在放大后的网格上重叠覆盖不均匀，容易导致生成的图像出现类似棋盘的微小网格斑纹。
- 替代方案：在现代工程中，常使用
nn.Upsample(mode='bilinear') + nn.Conv2d的组合来代替ConvTranspose2d，这种组合既能平滑放大图像，又能完全避免棋盘效应。
5.2.4 利用模型块组装U-Net
利用上述4个模型块，即可大体上完成拼接出UNet模型的实现。
以下代码实现来自unet/unet_model.py文件，笔者同样给其添加上了详细的注释，注明了经过每一个模型块前中后数据的变化。
1 | |
可以看到使用封装后的4个模型块来实现UNet模型时，代码非常优雅、代码量很少、可读性很好。
5.2.5 执行效果
准备工作： 1
2
3
4
5
6(bash) root@controller01:~# conda activate py31213
(py31213) root@controller01:~# cd /opt/code_repos/Pytorch-UNet
# 安装依赖的python库
(py31213) root@controller01:/opt/code_repos/Pytorch-UNet# pip install -r requirements.txt
# 下载训练时所需要的数据（从kaggle上下载）
(py31213) root@controller01:/opt/code_repos/Pytorch-UNet# bash scripts/download_data.sh
执行训练： 1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26"""
对于NVIDIA A40环境，训练时默认不要启用amp参数。
对于 A40 服务器环境中的官方 `train.py`，不要直接使用 `--amp`。该参数会启用 CUDA AMP 默认的 FP16 混合精度；训练中会导致数值不稳定，loss 出现 `Inf/NaN`，从而导致模型无法正常学习，最终生成无效权重，甚至导致恶性循环。可能的异常点包括模型前向计算、`softmax`、loss 或反向传播中的梯度计算。
"""
# 这里只启动一个轮次。默认是5 epochs
(py31213) root@controller01:/opt/code_repos/Pytorch-UNet# python train.py --epochs 1
...
INFO: Starting training:
Epochs: 1
Batch size: 1
Learning rate: 1e-05
Training size: 4580
Validation size: 508
Checkpoints: True
Device: cuda
Images scaling: 0.5
Mixed Precision: False
/opt/code_repos/Pytorch-UNet/train.py:80: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.
grad_scaler = torch.cuda.amp.GradScaler(enabled=amp)
Epoch 1/1: 20%|█████████████████████████████████████ | 916/4580 [01:57<07:37, 8.01img/s, loss (batch)=0.113]INFO: Validation Dice score: 0.8989954590797424 ...
Epoch 1/1: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4580/4580 [11:48<00:00, 6.46img/s, loss (batch)=0.0236]
INFO: Checkpoint 1 saved!
wandb:
wandb: 🚀 View run restful-yogurt-4 at: https://wandb.ai/{W&B 账号/用户}/U-Net/runs/77uzryk5
wandb: Find logs at: wandb/run-20260915_173344-77uzryk5/logs
执行预测： 1
2
3
4
5
6
7
8
9(py31213) root@controller01:/opt/code_repos/Pytorch-UNet# python predict.py \
> --model ./checkpoints/checkpoint_epoch1.pth \
> --input ./data/imgs/00087a6bd4dc_01.jpg \
> --output ./predict_output/output_fp32_epoch1.jpg
INFO: Loading model ./checkpoints/checkpoint_epoch1.pth
INFO: Using device cuda
INFO: Model loaded!
INFO: Predicting image ./data/imgs/00087a6bd4dc_01.jpg ...
INFO: Mask saved to ./predict_output/output_fp32_epoch1.jpgoutput_fp32_epoch1.jpg如下图所示： 
5.3 PyTorch修改模型
利用PyTorch，虽然我们可以从零开始构建模型，但有时这样从零开始可能会比较低效，尤其是我们想做快速验证时。此时我们可以挑选一些使用场景跟自己类似的开源模型，以其为基础进行修改、二开，修改原开源模型的结构或输入或输出，使其满足我们的需求。
此章节讲述如何修改基于PyTorch的模型，讲述的内容包括：修改模型层、添加额外输入、添加额外输出。 不管是如何修改原模型，首先要掌握原模型的使用场景、优缺点与结构和前向传播过程。在确定了原模型后，熟练了解原模型结构和前向传播过程对于自定义模型是非常有必要的。
5.3.1 修改模型层
本节以官方视觉库torchvision中内置的ResNet-50模型为基础，讲述如何修改模型层。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40import torch.nn as nn
from collections import OrderedDict
import torchvision.models as models
net = models.resnet50()
# 打印官方原 resnet50 模型的结构，因为内容比较多，这里只打印其 fc 层
print("官方原 resnet50 模型中的fc层结构：")
print(net.fc)
# 以下开始在原resnet-50的基础上修改其模型结构
# 使用前面学习过的nn.Sequential与collections.OrderedDict 自定义模型层并给其中子层设置Key
classifier = nn.Sequential(OrderedDict(
[('fc1', nn.Linear(2048, 128)), # 全连接层（2048个输入特征->128个输出特征）
('relu1', nn.ReLU()), # ReLU激活层
('dropout1',nn.Dropout(0.5)), # Dropout层
('fc2', nn.Linear(128, 10)), # 再一个全连接层（128个输入特征->10个输出特征）
# ('output', nn.Softmax(dim=1)) # Softmax层，将原始得分转成概率分布
]))
net.fc = classifier
# 修改 resnet50 模型结构后，再次打印其结构。这里只打印其中修改后 fc 层
print("\n修改了resnet50模型的fc层，再次打印其fc层结构：")
print(net.fc)
"""
# 执行时输出结果如下：
官方原 resnet50 模型中的fc层结构：
Linear(in_features=2048, out_features=1000, bias=True)
修改了resnet50模型的fc层，再次打印其fc层结构：
Sequential(
(fc1): Linear(in_features=2048, out_features=128, bias=True)
(relu1): ReLU()
(dropout1): Dropout(p=0.5, inplace=False)
(fc2): Linear(in_features=128, out_features=10, bias=True)
# (output): Softmax(dim=1)
)
"""
在resnet-50现有官方实现的基础上，我想对其进行少量的改造使其能够做一个10分类的问题。最简单的方法就是对其现有的fc层进行改造。我们不使用的一层全连接操作，因为它无法解决稍微复杂点的问题（如异或问题），暂时也没有必要直接堆叠多个全连接层，所以此处暂时保采用两个全连接层，中间包含激活层与Dropout层。
1
2
3
4
5
6
7
8classifier = nn.Sequential(OrderedDict(
[('fc1', nn.Linear(2048, 128)), # 全连接层（2048个输入特征->128个输出特征）
('relu1', nn.ReLU()), # ReLU激活层
('dropout1',nn.Dropout(0.5)), # Dropout层
('fc2', nn.Linear(128, 10)) # 再一个全连接层（128个输入特征->10个输出特征）
]))。
net.fc = classifiernn.Sequential与collections.OrderedDict将一系列简单的层组合起来，将最终的nn.Sequential实例直接赋值给net.fc即可实现上述改造目的。
5.3.2 添加额外输入
在基于开源或已有模型的基础上，来改造、训练自己的新模型时，原有模型的原有训练数据可能还能利用，但需要再添加一些额外的信息，此时在原有模型的工作流程过程中插入新模型的相关数据与流程。 为什么还要继续使用原有训练数据呢？因为成本、必要性等因素不想直接对原有训练数据进行再构建。此时我们可以在原模型前向传播到某个步骤时，添加额外的输入甚至与原模型当时的中间产生数据进行融合，并根据情况自定义后续的处理流程。
基本的思路是：将原模型添加额外输入前的部分作为一个整体继承或保留到新模型中，接收额外输入，根据新模型的需求构建额外输入与剩下流程之间的衔接与定义。
一般来说，新模型是通过继承torch.nn.Module来自定义，将原模型实例当成一个参数传递给新模型的__init__方法并对此方法中定义好新模型中需要的各模块与层，同时新模型的forward方法会接收那个额外输入后将其应用于新模型的前向传播过程中（再次说明需要对原模型的结构与前向传播过程很熟悉）。
如下以ResNet-50为原模型为基础，添加额外输入，定义了一个用于10分类的新模型（原ResNet-50是1000分类模型）。以下是可直接执行的源码实现。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85import torch
import torch.nn as nn
import torchvision.models as models
class Model(nn.Module):
"""在 ResNet-50 原 fc 层之前融合外部特征的 10 分类模型。"""
def __init__(self, net, add_feature_dim=1, num_classes=10):
super(Model, self).__init__()
# 1. 原模型中外部输入数据应用前的那部分结构及相关参数
# ResNet-50 的 children() 最后两个模块依次为 avgpool、fc。
# 这里保留 avgpool、去掉 fc，因此 image_features 的输出为
# (N, 2048, 1, 1)，也就是原最终全连接层之前的图像特征。
self.image_features = nn.Sequential(*list(net.children())[:-1])
image_feature_dim = net.fc.in_features # ResNet-50 中为 2048
# 2. 新模型中定义的 Dropout 层
self.dropout = nn.Dropout(p=0.5)
# 3. 新模型中定义的 全连接 层
# 图像特征 (N, 2048) 与外部特征 (N, add_feature_dim) 拼接后分类。
self.classifier = nn.Linear(
image_feature_dim + add_feature_dim, num_classes
)
def forward(self, images, add_variable):
# 1. 这里就是：原模型添加输入位置前的部分作为一个整体继续以其原工作流程工作，进行前向传播与展平
# 不会执行原 ResNet-50 的 net.fc；只提取其倒数第二层的特征。
image_features = self.image_features(images)
image_features = torch.flatten(image_features, start_dim=1) # (N, 2048)
# 2. 添加的外部输入、外部输入与原模型前置部分的输出的糅合
# 传入一维张量 (N,) ，进行形状对齐，统一成 (N, 1)。
if add_variable.ndim == 1:
add_variable = add_variable.unsqueeze(1)
# 此时如果add_variable 不是二维张量，后续的拼接肯定会出错
if add_variable.ndim != 2:
raise ValueError(
"add_variable 必须是 (batch_size, add_feature_dim) 的二维特征。"
)
if add_variable.shape[0] != image_features.shape[0]:
raise ValueError("images 与 add_variable 的 batch_size 必须一致。")
# torch.cat 后为 (N, 2048 + add_feature_dim)，即在最终 fc 前完成融合。
fused_features = torch.cat(
(self.dropout(image_features), add_variable), dim=1
)
# 3. 将糅合后的数据应用于新模型新建的全连接层上
return self.classifier(fused_features) # (N, 10)，供 CrossEntropyLoss 使用的 logits
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = models.resnet50()
# 额外变量将使得原模型输出的中间张量对应的维度上特征个数增加 1
add_feature_dim = 1
model = Model(net, add_feature_dim=add_feature_dim, num_classes=10).to(device)
model.eval() # 演示推理；训练时改为 model.train()
batch_size = 3
inputs = torch.randn(batch_size, 3, 224, 224, device=device)
# 每张图一个标量的额外信息，例如年龄、传感器读数等。
add_var = torch.randn(batch_size, add_feature_dim, device=device)
with torch.no_grad():
outputs = model(inputs, add_var)
print("运行设备：", device)
print("图像输入形状：", inputs.shape)
print("外部特征形状：", add_var.shape)
print("分类 logits 形状：", outputs.shape) # torch.Size([3, 10])
"""
# 执行时输出结果：
运行设备： cpu
图像输入形状： torch.Size([3, 3, 224, 224])
外部特征形状： torch.Size([3, 1])
分类 logits 形状： torch.Size([3, 10])
"""
可以看到，自定义模型Model的__init__方法接收了net、add_feature_dim、num_classes，它们分别是ResNet-50实例、添加的额外输入的某方面信息（此处为其要增加的特征个数）、新模型的最终分类数。此外，此方法中将原ResNet-50中的fc层之前的结构赋值给了属性image_features以将它们全部保留了下来。定义了新的
Dropout层 与
全连接层，其中全连接层是适配了新的输入特征数与输出特征数的全连接层。
在自定义模型Model的forward方法中除了接收原输入数据images外，还接收了额外输入add_variable。forward方法中首先将images应用于原ResNet-50中的fc层之前的结构进行前向传播，并将此前向传播的结果数据进行展平（原ResNet-50模型中也是在其forward中临时定义的，新模型中无法继承也只能仿照着，自己来定义），然后将展平得到的数据进行部分舍弃后与额外输入数据进行融合（融合并对额外输入进行了结构等检查），此时最后的fused_features的形状是(N,
2048 +
add_feature_dim)即(N, 2049)，最终通过全连接层运算后，得到N个样本在10分类上的logits（原始得分）。
5.3.3 添加额外输出
基于已有原模型进行二开改造得到新模型后，我们一般主要关注新模型的最后输出，但有时我们也想关注原模型某部分的输出结果或新模型的某中间层输出结果，用于在训练循环中用于计算辅助损失、获得及时的反馈监督，以便更好地调整中间层参数等、获得更好的模型训练效果。
常规的做法是在新模型的forward方法中获取到我们关注的中间结果，并把它（们）与新模型的主结果一起返回（python中方法支持返回多个结果）。
仍然以ResNet-10为原模型，在上个示例的基础上稍加修改得到如下示例。其中展示了新模型返回”原
ResNet fc 前的中间特征“与”主分类结果“。 1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106import torch
import torch.nn as nn
import torchvision.models as models
class Model(nn.Module):
"""返回 ResNet 中间特征，并用融合外部特征后的分类头完成 10 分类。
``forward`` 同时返回：
1. ``image_features``：原 ResNet-50 的 ``fc`` 层之前的 2048 维特征；
训练循环可取得它并据此计算辅助损失，向中间层施加额外监督。
2. ``logits``：融合外部特征后得到的 10 类主任务预测结果。
注意：本类负责产生两个结果；辅助损失应在训练循环中根据具体的
监督目标计算，而不在 ``forward`` 中直接计算。
"""
def __init__(self, net, add_feature_dim=1, num_classes=10):
super(Model, self).__init__()
# 1. 原模型中、添加外部输入之前的部分。
# ResNet-50 的 children() 最后两个模块依次为 avgpool、fc。
# 此处保留 avgpool、去掉原 fc：因此不会获得原模型的 1000 类输出，
# 而是取得原 fc 前的图像特征 (N, 2048, 1, 1)。
self.image_features = nn.Sequential(*list(net.children())[:-1])
image_feature_dim = net.fc.in_features # ResNet-50 中为 2048
# 2. 新模型的主分类头：用于融合特征后的 10 分类。
self.dropout = nn.Dropout(p=0.5)
# 图像特征 (N, 2048) 与外部特征 (N, add_feature_dim) 拼接后分类。
self.classifier = nn.Linear(
image_feature_dim + add_feature_dim, num_classes
)
def forward(self, images, add_variable):
# 1. 图像经过原 ResNet-50 的主体。原 net.fc 不会执行。
image_features = self.image_features(images) # (N, 2048, 1, 1)
image_features = torch.flatten(image_features, start_dim=1) # (N, 2048)
# 这个中间特征被保留到 return 中，供训练循环施加辅助监督。
# 2. 对齐外部输入形状，再与图像特征融合。
# 若每张图只有一个标量外部特征，可传入 (N,)；此处将其变为 (N, 1)。
if add_variable.ndim == 1:
add_variable = add_variable.unsqueeze(1)
# 形状对齐之后，外部特征必须为 (N, add_feature_dim)。
if add_variable.ndim != 2:
raise ValueError(
"add_variable 必须是 (batch_size, add_feature_dim) 的二维特征。"
)
if add_variable.shape[0] != image_features.shape[0]:
raise ValueError("images 与 add_variable 的 batch_size 必须一致。")
# 拼接结果为 (N, 2048 + add_feature_dim)。
fused_features = torch.cat(
(self.dropout(image_features), add_variable), dim=1
)
# 3. 同时返回中间特征和主任务结果。
# image_features: (N, 2048)，即原 ResNet fc 前的中间特征；
# 可在训练循环中用于计算辅助损失。
# logits: (N, 10)，主分类结果；未经过 Softmax，可直接传给 CrossEntropyLoss。
logits = self.classifier(fused_features)
return image_features, logits
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = models.resnet50()
# 每张图融合一个标量外部特征，因此主分类头输入为 2048 + 1 = 2049 维。
add_feature_dim = 1
model = Model(net, add_feature_dim=add_feature_dim, num_classes=10).to(device)
model.eval() # 演示推理；训练时改为 model.train()
batch_size = 3
inputs = torch.randn(batch_size, 3, 224, 224, device=device)
# 每张图一个标量的额外信息，例如年龄、传感器读数等。
add_var = torch.randn(batch_size, add_feature_dim, device=device)
with torch.no_grad():
# 训练时不应使用 no_grad()，因为主损失和辅助损失都需要反向传播。
image_features, logits = model(inputs, add_var)
print("运行设备：", device)
print("图像输入形状：", inputs.shape)
print("外部特征形状：", add_var.shape)
print("ResNet 原 fc 前特征形状：", image_features.shape) # [3, 2048]
print("分类 logits 形状：", logits.shape) # [3, 10]
# 训练时的概念示例（不在本示例中执行）：
# main_loss = nn.CrossEntropyLoss()(logits, class_labels)
# auxiliary_loss = ... # 根据额外监督目标，由 image_features 计算
# total_loss = main_loss + auxiliary_weight * auxiliary_loss
"""
# 执行时输出结果：
运行设备： cpu
图像输入形状： torch.Size([3, 3, 224, 224])
外部特征形状： torch.Size([3, 1])
ResNet 原 fc 前特征形状： torch.Size([3, 2048])
分类 logits 形状： torch.Size([3, 10])
"""
5.4 PyTorch模型保存与读取
基于PyTorch构建好模型，训练后达到一定效果后，可以将其保存下来，以便后续直接恢复训练得到的模型及效果。 本章节学习如何进行PyTorch模型的保存与读取。以下都是基于PyTorch这一机器学习与深度学习框架而言，具体包含的内容有：模型的存储格式、模型的存储内容、单卡或多卡下模型的存储与加载、模型权重外其他参数的存取。
5.4.1 模型存储内容
PyTorch 模型主要包含两个部分：模型结构和模型权重（状态参数）。保存
PyTorch
模型时，通常有两类方式：保存整个模型架构，或仅保存模型权重。这两者都是通过
torch.save()
方法来实现的，区别仅在于传入的保存对象不同。
在 PyTorch 中，自定义模型通常继承自
torch.nn.Module。模型结构是由在 __init__
中定义的各个子模块或层（如卷积层、全连接层）所构建的计算图网络；而模型权重则是指这些子模块或层中的参数张量（Tensor），主要保存在模型的
state_dict（状态字典）中。这些权重既包含在训练过程中通过反向传播和优化器更新的可训练参数（Parameters，如权重与偏置），也包含在前向传播中动态更新的缓冲状态（Buffers，如
BatchNorm 的运行均值与方差）。
1 | |
这个程序执行会在本地生成如下两个文件，它们的大小几乎是一样的。
这说明：对 ResNet-50
这类参数量很大的模型，模型文件空间的绝大多数由权重/缓冲区张量占用；完整保存模型比仅保存
state_dict() 多出的结构序列化信息相对很小。
1
2
3
4
5
6
7
8
9
10
11ai_py312) jsy@jmac 05第五章 % ls -al *.pth
-rw-r--r--@ 1 jsy staff 102544474 9月 18 11:11 resnet-50-weights.pth
-rw-r--r--@ 1 jsy staff 102569018 9月 18 11:11 resnet-50.pth
(ai_py312) jsy@jmac 05第五章 % ls -alh *.pth
-rw-r--r--@ 1 jsy staff 98M 9月 18 11:11 resnet-50-weights.pth
-rw-r--r--@ 1 jsy staff 98M 9月 18 11:11 resnet-50.pth
(ai_py312) jsy@jmac 05第五章 % md5 resnet-50.pth resnet-50-weights.pth
MD5 (resnet-50.pth) = ab9c074db3f874cf4402e073238f9c67
MD5 (resnet-50-weights.pth) = a25bbbbdd219466c359ae4eba6a6d09b
5.4.2 模型存储格式
上述示例中是以.pth的格式来保存整个PyTorch模型或模型权重的，但其实PyTorch中还支持以pt或pkl
这两种格式来保存整个PyTorch模型或模型权重。仅从功能与实用性来说，它们的差别可以忽略，就是随便采用哪种。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21import torch
from torchvision import models
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
save_dict = {
1: "pth",
2: "pt",
3: "pkl"
}
save_preffix = "resnet-50"
for i in range(1, 4):
print(f"模型或模型权重以 {save_dict[i]} 格式保存...")
save_whole = f'./{save_preffix}.{save_dict[i]}'
save_weights = f'./{save_preffix}-weights.{save_dict[i]}'
# 保存整个模型
torch.save(model, save_whole)
# 只保存模型权重
torch.save(model.state_dict(), save_weights)
上述示例代码执行后在本地生成6个文件，它们都是针对resnet-50模型有效的存储。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15(ai_py312) jsy@jmac 05第五章 % ls -al *.pt *.pth *.pkl
-rw-r--r--@ 1 jsy staff 102544474 9月 18 11:16 resnet-50-weights.pkl
-rw-r--r--@ 1 jsy staff 102544474 9月 18 11:16 resnet-50-weights.pt
-rw-r--r--@ 1 jsy staff 102544474 9月 18 11:16 resnet-50-weights.pth
-rw-r--r--@ 1 jsy staff 102569018 9月 18 11:16 resnet-50.pkl
-rw-r--r--@ 1 jsy staff 102569018 9月 18 11:16 resnet-50.pt
-rw-r--r--@ 1 jsy staff 102569018 9月 18 11:16 resnet-50.pth
(ai_py312) jsy@jmac 05第五章 % cat md5_results.txt
ab9c074db3f874cf4402e073238f9c67 ./resnet-50.pth
893f71d417b23f2cb3c7010012911b11 ./resnet-50-weights.pth
ab9c074db3f874cf4402e073238f9c67 ./resnet-50.pt
893f71d417b23f2cb3c7010012911b11 ./resnet-50-weights.pt
ab9c074db3f874cf4402e073238f9c67 ./resnet-50.pkl
893f71d417b23f2cb3c7010012911b11 ./resnet-50-weights.pkl
这说明：torch.save()
根据“被保存的对象”和其序列化规则写入内容，不会因为路径后缀是
.pth、.pt 或 .pkl
而切换格式。加载时也不需要“后缀匹配”，例如这个文件即便改名为
.abc即使用torch.load("resnet-50.abc", weights_only=False)，仍可用。
实践中一般选一种统一命名（常见为 .pth 或
.pt），并优先保存
state_dict()；后缀本身不提供额外兼容性或压缩效果。
5.4.3 单卡和多卡模型存储的区别
本章节代码在如下环境中验证通过： 1
2
3
4
5容器操作系统：Ubuntu 22.04
PyTorch版本： 2.5.1+cu124
Python版本： 3.12
NVIDIA GPU：RTX 2080 Ti x2，其中每个GPU有 22 GB 显存。
CUDA版本： 12.4
在单卡训练时，model 通常就是原始的
nn.Module 对象；其 state_dict() 中的参数名类似
conv1.weight、layer1.0.conv1.weight。使用
nn.DataParallel（DP）或
nn.parallel.DistributedDataParallel（DDP）进行多卡训练时，原始模型会被包装到包装器的
module 属性中：
1 | |
因此，若直接保存多卡包装器的
model.state_dict()，参数键名会多出 module.
前缀，例如 module.conv1.weight。而单卡模型的
state_dict() 没有此前缀。将前者（即带有
module.
前缀的多卡包装器权重）直接加载到未包装的单卡模型中时，load_state_dict()
会因参数名不匹配而报出 Missing key(s) 和
Unexpected key(s)。
下面是可在双 GPU 环境中直接执行的 DP 示例。使用
weights=None，因此不会下载预训练权重；本示例仅演示包装、键名和存储，不进行训练。
1 | |
为避免检查点与“当前到底使用几张卡”绑定，推荐始终保存原始模型的权重：单卡模型保存
model.state_dict()；DP/DDP 多卡模型保存
model.module.state_dict()。这样两种情况得到的键名一致，均可加载到未包装的模型中。
上例中 resnet50_dp_weights.pth
就是推荐的通用权重文件。若同一份代码既可能在单卡运行，也可能在 DP/DDP
下运行，可使用下面的完整保存代码：
1 | |
读取权重时还应显式指定 map_location。例如在 CPU
或不同编号的 GPU 上读取检查点时，先映射到
CPU，再将模型移动到目标设备，可以避免检查点中记录的原设备不可用的问题。下面的完整示例先保存一份权重，再在目标设备上读取它：
1 | |
这里保存的是
state_dict，因此不包含模型结构；加载前仍需要用相同的模型定义先创建
model。相比 torch.save(model, path)
直接保存整个 Python 对象，这种方式对类的导入路径和代码环境依赖更少，也是
PyTorch 更推荐的模型权重保存方式。
需要注意，DP 与 DDP 的训练机制不同：DP 在一个进程中管理多张卡，DDP
通常采用“一张 GPU 对应一个进程”。但二者都通过 .module
暴露被包装的原始模型，所以本节关于保存
model.module.state_dict() 和 module.
前缀的结论对二者都适用。DDP 训练时各进程的模型权重保持同步，通常只让
rank 0 进程写入检查点，避免多个进程同时写同一文件。DDP
的完整示例还必须通过 torchrun
初始化进程组，本节先以可直接运行的 DP 示例说明模型存储规则。
5.4.4 单卡/多卡情况分类讨论
本章节代码在如下环境中验证通过： 1
2
3
4
5容器操作系统：Ubuntu 22.04
PyTorch版本： 2.5.1+cu124
Python版本： 3.12
NVIDIA GPU：RTX 2080 Ti x2，其中每个GPU有 22 GB 显存。
CUDA版本： 12.4
按照上一节的推荐做法——无论训练时是否使用 DP/DDP，都保存原始模型的
state_dict()——单卡与多卡之间的读取不再需要修改参数名。下面以
models.resnet50(weights=None)
为例；若是自定义模型，创建时也必须使用与保存时相同的网络结构和超参数。
5.4.4.1 单卡保存，单卡或 CPU 加载
单卡训练时直接保存模型权重。读取时先创建未包装的模型，加载权重后再移动到运行设备。
1 | |
5.4.4.2 单卡保存，多卡加载
仍然先向未包装模型加载权重，再按照当前运行环境包装为 DP 或 DDP。权重文件本身不需要改变。
大致流程就是： 1
2
3
4
5
6
7
8
9单卡训练模型
→ 保存 model.state_dict()
多卡环境
→ 用相同模型结构创建未包装的模型实例loaded_model
→ 通过 map_location="cpu" 读取权重张量为state_dict
→ loaded_model.load_state_dict(state_dict)
→ 将原始模型移动到主卡 cuda:0
→ 按当前环境包装为 DP 或 DDP
1 | |
5.4.4.3 多卡保存，单卡或 CPU 加载
多卡训练结束时，保存包装器内原始模型的权重，而不是保存
model.state_dict()。读取时与单卡保存的情况完全相同。
大致流程就是： 1
2
3
4
5
6
7
8
9
10
11多卡环境
→ 创建或恢复原始模型后，移动到主卡 cuda:0
→ 按当前环境包装为 DP 或 DDP
→ [可能进行训练操作]
→ 以”dp_model.module.state_dict()“形式只保存模型的权重
单卡环境
→ 用相同模型结构创建模型实例
→ 通过 map_location="cpu" 读取权重张量为state_dict
→ loaded_model.load_state_dict(state_dict)
→ 将原始模型移动到cuda:0
1 | |
5.4.4.4 多卡保存，多卡加载
保存和加载都遵循“保存、加载原始模型；需要训练时再包装”的顺序即可。下面的完整示例使用
DP。若改用 DDP，还应确保仅 rank 0 保存，其他进程在需要使用新检查点前通过
torch.distributed.barrier() 等待文件写入完成。
大致流程就是： 1
2
3
4
5
6
7
8
9
10
11
12多卡环境1
→ 创建或恢复原始模型后，移动到主卡 cuda:0
→ 按当前环境包装为 DP 或 DDP
→ [可能进行训练操作]
→ 以”dp_model.module.state_dict()“形式只保存模型的权重
多卡环境2
→ 用相同模型结构创建模型实例
→ 通过 map_location="cpu" 读取权重张量为state_dict
→ loaded_model.load_state_dict(state_dict)
→ 将原始模型移动到主卡 cuda:0
→ 按当前环境包装为 DP 或 DDP
1 | |
5.4.4.5 兼容已经保存的
module. 前缀检查点
如果历史代码使用了 torch.save(model.state_dict(), path)
保存 DP/DDP 包装器，文件中的键名已经带有 module.
前缀。加载到未包装模型前，可删除每个键名开头的这一个前缀：
关键代码在于： 1
2
3
4
5
6
7
8
9
10
11
12
13
14
15# 此时，state_dict 是恢复 多卡场景下通过如下形式保存下来的模型权重 得到的：
torch.save(dp_model.state_dict(), legacy_path)
# 先将保存时位于 GPU 的张量映射到 CPU；此时键名仍保留 module. 前缀。
state_dict = torch.load(
legacy_path,
map_location="cpu",
weights_only=True)
# 字典推导式创建新状态字典：只删除每个键开头的一个 module. 前缀。
# removeprefix() 不会误删键名中间的相同字符串，比 replace() 更精确。
state_dict = {
key.removeprefix("module."): value
for key, value in state_dict.items()
}
完整代码如下： 1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43# 50-加载带module前缀的权重.py
import torch
import torch.nn as nn
from torchvision import models
# 本示例先构造旧式 DP 检查点，因此必须至少有两张可见 GPU。
if torch.cuda.device_count() < 2:
raise RuntimeError("本示例需要至少两张可见的 NVIDIA GPU。")
# legacy 表示历史遗留格式：其中参数键名带有 module. 前缀。
legacy_path = "resnet50_legacy_dp.pth"
# 模拟历史代码：dp_model 是 DP 包装器，直接调用 dp_model.state_dict()
# 会将包装器名称写入键名，例如 module.conv1.weight。
# 这类文件不能直接加载到未包装的 ResNet-50，需要先处理键名。
dp_model = models.resnet50(weights=None).to("cuda:0")
dp_model = nn.DataParallel(dp_model, device_ids=[0, 1])
torch.save(dp_model.state_dict(), legacy_path)
# 先将保存时位于 GPU 的张量映射到 CPU；此时键名仍保留 module. 前缀。
state_dict = torch.load(
legacy_path,
map_location="cpu",
weights_only=True)
# 字典推导式创建新状态字典：只删除每个键开头的一个 module. 前缀。
# removeprefix() 不会误删键名中间的相同字符串，比 replace() 更精确。
state_dict = {
key.removeprefix("module."): value
for key, value in state_dict.items()
}
# 处理后的键名已与普通 ResNet-50 一致，所以此处创建未包装模型。
loaded_model = models.resnet50(weights=None)
loaded_model.load_state_dict(state_dict)
# 成功加载后可直接 CPU 推理；若需要 GPU 推理，再执行 loaded_model.to("cuda:0")。
loaded_model.eval()
print("历史 DP 检查点已成功加载到未包装模型。")
"""
### 执行时的输出：
历史 DP 检查点已成功加载到未包装模型。
"""
也可以先创建 DP/DDP 包装器，再将带 module.
前缀的权重加载到包装器中；但这会使加载代码依赖当前的并行方式。对于历史检查点，去除前缀后加载到原始模型更通用；对于新检查点，直接保存
model.module.state_dict() 可以从源头避免此问题。
5.4.5 其他参数的保存和读取
在实际训练中，若只保存模型的
state_dict()，虽然可以用于推理或微调，却不能准确地从中断处继续训练。因为优化器（如
Adam 的动量状态）、学习率调度器、已完成的
epoch、当前损失或最佳验证指标等训练状态都会丢失。
为了能够恢复训练，通常将这些内容组成一个字典并保存为训练检查点（checkpoint）：
1 | |
在 PyTorch 2.6.1 中，torch.load() 的
weights_only 默认值为
True。对于本节这种仅由张量、整数、浮点数、字符串、列表和普通字典组成的检查点，应显式写出
weights_only=True。若检查点包含自定义 Python
类等任意对象，才可能需要 weights_only=False；这会执行
pickle 反序列化，因此只能读取可信来源的文件。
下面是一个可独立运行的完整示例。它使用随机数据完成一次训练更新，保存训练检查点，再新建模型、优化器和调度器以模拟重新启动程序后恢复训练。device
会在 GPU 可用时使用 cuda:0，否则使用
CPU；weights=None 不会下载预训练权重。
1 | |
该程序的 loss
数值会因随机初始化和随机输入而变化，但其输出结构应如下：
1 | |
其中，scheduler.step() 已将学习率从 0.01
调整为
0.001。如果只恢复模型权重、不恢复优化器和调度器状态，下一轮训练会错误地从新建优化器的学习率
0.01 开始，且 SGD 动量等历史状态也会丢失。
对于 DP/DDP 训练，模型权重仍遵循前文的规则：保存
model.module.state_dict()；优化器和调度器的
state_dict() 保存方式不变。使用 DDP 时，通常只允许 rank 0
进程写检查点，并在其他进程读取前进行同步。除训练检查点外，如果还需保证恢复后与中断前的数据顺序和随机行为完全一致，还应额外保存并恢复随机数生成器状态、数据采样器状态等信息。