文章目录
- 1 模型构造
- 1.1 自定义 MLP(多层感知机)
- 1.1.1 `__init__` (构造函数)
- 1.1.2 `forward` (前向传播)
- 1.2 使用自定义 MLP
- 1.3 自定义 `Sequential` 类
- 1.4 前向传播
- 1.5 模块的嵌套使用
- 2 参数管理
- 2.1 参数访问
- 2.2 嵌套模型
- 2.3 参数初始化
- 2.4 参数共享
- 3 自定义层
- 3.1 自定义无参数层
- 3.2 自定义带参数层
- 3.3 自定义层的应用
- 4 读写文件
- 4.1 保存与加载基本张量
- 4.1.1 保存对象
- 4.1.2 加载对象
- 4.2 保存与加载模型
- 4.2.1 保存模型参数
- 4.2.2 加载模型参数
- 5 使用 GPU
- 5.1 确认和查看 GPU 信息
- 5.2 指定计算设备 (Device)
- 5.3 张量 (Tensor) 与 GPU
- 5.3.1 创建 GPU 上的张量
- 5.3.2 GPU 上的张量运算
- 5.4 神经网络与 GPU
1 模型构造
在 PyTorch 中,nn.Module
是所有神经网络模块(包括层和整个模型)的基类。通过继承它,我们可以构建更加灵活和复杂的网络模型。
nn.Module
: PyTorch 中所有网络层和模型的基类。nn.functional
(通常简写为F
): 包含许多没有可训练参数的函数,例如激活函数(如ReLU
)和池化操作。nn.Sequential
: 一个特殊的nn.Module
子类,可以按顺序封装多个层,提供一种简单的线性模型构建方式。
1.1 自定义 MLP(多层感知机)
要自定义一个多层感知机,我们需要继承 nn.Module
并实现两个关键方法:__init__
和 forward
。
1.1.1 __init__
(构造函数)
这个方法用于定义模型所需的各个层(或子模块)。
- 首先,需要调用父类的构造函数
super().__init__()
来进行必要的初始化。 - 接着,定义你需要的各种层,并将它们赋值给类的成员变量,例如
self.hidden
和self.out
。 - 注意:在
__init__
中定义的层会被 PyTorch 自动识别为模型的一部分,它们的参数也会被自动追踪。
import torch
from torch import nnclass MLP(nn.Module):def __init__(self):super().__init__()# 定义隐藏层,输入维度20,输出256self.hidden = nn.Linear(20, 256)# 定义输出层,输入维度256,输出10self.out = nn.Linear(256, 10)
1.1.2 forward
(前向传播)
这个方法定义了数据如何流过这些层,也就是计算逻辑。
- 它接收输入数据
x
作为参数。 - 通过调用在
__init__
中定义的层,并使用nn.functional
中的函数,一步步完成前向计算。 - 最后,返回计算结果。
from torch.nn import functional as Fclass MLP(nn.Module):# ... __init__ 方法 ...def forward(self, x):# 将输入x通过隐藏层和ReLU激活函数x = F.relu(self.hidden(x))# 将结果通过输出层return self.out(x)

1.2 使用自定义 MLP
使用自定义的 MLP
类和使用 nn.Sequential
的方式类似,先实例化类,然后传入数据。
# 实例化 MLP
mlp = MLP()
# 创建一个随机输入数据
X = torch.rand(2, 20)
# 得到输出
output = mlp(X) # 输出的形状将是 (2, 10)

1.3 自定义 Sequential
类
李沐老师展示了如何通过继承 nn.Module
来实现一个功能与 nn.Sequential
类似的自定义类 MySequential
。
- 在
__init__
中,它接收一个可变参数*args
(表示一系列层)。 - 它使用
nn.ModuleList
或者nn.ModuleDict
等容器来存储这些层。 - 在
forward
中,它通过循环遍历存储的层,按顺序将数据依次传递给每一层。
from collections import OrderedDictclass MySequential(nn.Module):def __init__(self, *args):super().__init__()# PyTorch 推荐使用这种方式存储子模块for idx, module in enumerate(args):self._modules[str(idx)] = moduledef forward(self, x):# 遍历所有子模块,按顺序进行前向计算for module in self._modules.values():x = module(x)return x
这个例子表明,nn.Sequential
并非一个魔法类,它的核心逻辑就是将层存储起来并按顺序调用,而这种逻辑完全可以由我们自己实现。

1.4 前向传播
继承 nn.Module
的最大优势在于,我们可以在 forward
方法中编写任意的计算逻辑,而不受 nn.Sequential
严格的线性流限制。
例如,可以在 forward
中加入:
- 非训练参数: 使用
torch.rand
创建一个不参与训练的随机权重random_weights
,并设置requires_grad=False
。 - 复杂的控制流: 例如
if
、for
或while
循环,根据数据值动态地调整计算过程。 - 多路径计算: 将输入数据同时传递给不同的层,再将结果进行合并。
- 自定义操作: 如矩阵乘法
torch.mm()
等。 - 返回标量: 不返回矩阵,而是返回一个求和后的标量值。
class FlexMLP(nn.Module):def __init__(self):super().__init__()self.linear = nn.Linear(20, 20)# 创建一个不参与训练的随机权重self.random_weights = torch.rand(20, 20, requires_grad=False)def forward(self, x):x = self.linear(x)# 在 forward 中加入自定义逻辑x = torch.mm(x, self.random_weights)x = x + 1x = F.relu(x)# 复杂的控制流while x.abs().sum() > 1:x /= 2# 返回一个标量return x.sum()
1.5 模块的嵌套使用
在 PyTorch 中,任何 nn.Module
的子类都可以作为另一个 nn.Module
的子模块,这使得我们能够非常灵活地构建复杂的网络结构。
例如,一个 Sequential
模块可以包含另一个自定义的 MLP
模块,而这个 MLP
模块内部又可以包含 Sequential
模块。这种嵌套和组合的能力让模型设计变得非常强大。

2 参数管理
我们定义如下网络。
import torch
from torch import nnnet = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 1))
X = torch.rand(size=(2, 4))
net(X)
2.1 参数访问
在 PyTorch 中,模型的参数(如权重和偏置)可以通过多种方式访问。
-
state_dict()
:- 这是一个有序字典(
OrderedDict
),存储了模型所有可学习参数的键值对。键是参数的名称,值是对应的 Tensor。 - 可以通过
net.state_dict()
访问整个模型的参数字典,也可以通过net[layer_index].state_dict()
访问特定层的参数。
- 这是一个有序字典(
-
net.parameters()
:- 返回一个迭代器,可以遍历模型中所有的 可优化参数(
Parameter
)。 Parameter
是Tensor
的子类,具有requires_grad=True
属性,并且会自动被注册为模型的一部分。- 可以通过
.data
属性访问参数的实际值,通过.grad
属性访问参数的梯度。
- 返回一个迭代器,可以遍历模型中所有的 可优化参数(
-
net.named_parameters()
:- 返回一个迭代器,它生成 (name, parameter) 的键值对,让你在访问参数的同时知道其在模型中的名称。
- 这个名称是基于模型结构自动生成的,例如
'0.weight'
、'0.bias'
、'2.weight'
等。
2.2 嵌套模型
当模型包含其他模型作为子模块时,PyTorch 会自动处理参数的嵌套关系。

- 打印模型结构:
- 直接
print(net)
可以以字符串形式直观地展示整个模型的嵌套结构,包括每个层的名称、类型和参数。
- 直接

2.3 参数初始化
PyTorch 提供了灵活的方式来初始化模型的参数,可以修改默认的初始化方法。
-
torch.nn.init
:- 这个模块包含了多种初始化函数,如
normal_()
(正态分布)、zeros_()
(全零)、constant_()
(常数)和xavier_uniform_()
(Xavier 均匀分布)等。
- 这个模块包含了多种初始化函数,如
-
apply()
方法:-
net.apply(init_func)
是一个非常强大的方法,它会递归地遍历模型的所有子模块,并对每个子模块调用指定的init_func
函数。 -
你可以在
init_func
中通过isinstance
或type
判断模块类型(如nn.Linear
),然后对符合条件的模块进行参数初始化。# 示例:自定义初始化函数 def init_normal(m):if type(m) == nn.Linear:nn.init.normal_(m.weight, mean=0, std=0.01)nn.init.zeros_(m.bias)# 对整个网络应用该初始化函数 net.apply(init_normal)
-
-
直接修改:
- 也可以直接访问参数的
.data
属性并进行修改。这种方法最直接,但可能不如apply()
灵活和安全。 - 例如:
net[0].weight.data.fill_(1.0)
将第一层的权重全部设置为 1。
- 也可以直接访问参数的

2.4 参数共享
在某些情况下,你可能希望不同的层共享同一组参数,这意味着它们的权重和偏置是完全相同的,并且在训练过程中会一起更新。
- 方法:
- 首先定义一个独立的层作为“共享层”,例如
shared_layer = nn.Linear(...)
。 - 然后在构建模型时,将这个共享层实例多次赋值给不同的子模块。
- PyTorch 会自动识别并将其视为同一个参数实例。
- 首先定义一个独立的层作为“共享层”,例如
# 示例:共享参数
# 首先定义一个共享层
shared = nn.Linear(8, 8)# 构建一个Sequential模型
net = nn.Sequential(nn.Linear(4, 8),shared, # 第二层nn.ReLU(),shared, # 第四层,与第二层共享参数nn.Linear(8, 1)
)# 此时,net[1].weight 和 net[3].weight 指向同一个对象
- 验证:
- 可以通过
is
运算符来判断两个参数是否为同一个对象实例。net[1].weight is net[3].weight
将返回True
。 - 修改其中一个参数的值,另一个参数的值也会随之改变,因为它们指向的是内存中的同一块数据。
- 可以通过

3 自定义层
3.1 自定义无参数层
自定义一个没有可训练参数的层非常简单,因为它本质上就是一个自定义的 nn.Module
。
- 定义类:
- 继承
torch.nn.Module
类。
- 继承
__init__
(构造函数):- 调用父类的构造函数
super().__init__()
。如果该层不需要参数,这一步可以省略(PyTorch 3+ 会自动添加),但为了代码清晰,建议保留。
- 调用父类的构造函数
forward
(前向传播):- 实现该层的计算逻辑。例如,李沐老师的例子中定义了一个
CenteredLayer
,其forward
方法的作用是将输入数据的均值减去,使其中心化。
- 实现该层的计算逻辑。例如,李沐老师的例子中定义了一个
import torch
from torch import nnclass CenteredLayer(nn.Module):def __init__(self):super().__init__()def forward(self, x):# 减去均值return x - x.mean()# 使用示例
layer = CenteredLayer()
input_tensor = torch.randn(5, 5)
output_tensor = layer(input_tensor)
print(output_tensor.mean()) # 输出值会非常接近0

3.2 自定义带参数层
如果自定义的层需要包含可学习的参数,你需要使用 torch.nn.Parameter
类。
torch.nn.Parameter
:- 这是
torch.Tensor
的一个子类,它的特殊之处在于,当它被赋值为nn.Module
的成员变量时,它会被自动添加到模型的参数列表中。 - 这意味着 PyTorch 的自动求导机制会追踪它的梯度,并在优化器进行参数更新时,自动对其进行优化。
- 这是
- 自定义带参数层的步骤:
- 在
__init__
中定义参数:- 使用
torch.randn
或其他初始化方法创建一个Tensor
作为参数的初始值。 - 将该
Tensor
包裹在nn.Parameter
中,并将其赋值给类的成员变量,例如self.weight
。
- 使用
- 在
forward
中使用参数:- 实现前向计算逻辑,将输入数据与
self.weight
和self.bias
等参数进行运算。 - 你可以直接访问
nn.Parameter
实例,因为它本身就是一个Tensor
,可以参与各种张量运算。
- 实现前向计算逻辑,将输入数据与
- 在
import torch
from torch import nnclass MyLinear(nn.Module):def __init__(self, in_features, out_features):super().__init__()self.in_features = in_featuresself.out_features = out_features# 使用 nn.Parameter 包装可学习参数self.weight = nn.Parameter(torch.rand(in_features, out_features))self.bias = nn.Parameter(torch.rand(out_features))def forward(self, x):# 使用 self.weight 和 self.bias 进行前向计算# 这里使用F.relu只是为了演示,可以根据需要替换return torch.matmul(x, self.weight) + self.bias# 或者使用torch.matmul(x, self.weight.data) + self.bias.data来访问值# 但通常直接使用参数本身即可# 使用示例
my_layer = MyLinear(4, 3)
print(my_layer.weight.shape)
print(my_layer.bias.requires_grad) # True, 因为是 nn.Parameter

3.3 自定义层的应用
- 自定义的层可以像 PyTorch 内置的层一样使用。
- 它们可以被放置在
nn.Sequential
容器中,与其他层一起构成更复杂的网络。 - 这种方式允许你将任何独特的计算逻辑封装成一个可复用的模块,极大地提高了代码的模块化和灵活性。

4 读写文件
4.1 保存与加载基本张量
PyTorch 的 torch.save
和 torch.load
函数是用于序列化和反序列化张量及其他 Python 对象的通用方法。
4.1.1 保存对象
torch.save(obj, filename)
:obj
: 你想要保存的 Python 对象,例如单个 Tensor、**列表(list)**或 字典(dictionary)。filename
: 保存文件的路径。
示例:
import torchx = torch.arange(4)
torch.save(x, 'x-file.pt')my_list = [x, torch.zeros(4)]
torch.save(my_list, 'my-list.pt')my_dict = {'x': x, 'y': torch.zeros(4)}
torch.save(my_dict, 'my-dict.pt')

4.1.2 加载对象
torch.load(filename)
:filename
: 要加载的文件的路径。- 该函数会返回保存时序列化的对象。
示例:
# 加载单个 Tensor
x2 = torch.load('x-file.pt')
# 加载列表
x3, y3 = torch.load('my-list.pt')
# 加载字典
my_dict2 = torch.load('my-dict.pt')

4.2 保存与加载模型
对于神经网络模型,我们通常只保存可学习的参数(权重和偏置),而不是整个模型定义。这是因为 PyTorch 的**命令式(imperative)**编程风格更侧重于计算图的动态构建,因此模型的结构定义需要由代码本身提供。
4.2.1 保存模型参数
state_dict()
:- 这是
nn.Module
的一个方法,它返回一个 Python 字典,其中包含了模型的所有参数。 - 键是参数的名称(例如
'0.weight'
),值是对应的Tensor
。
- 这是
- 保存步骤:
- 实例化一个模型,例如
net = MLP(...)
。 - 获取模型的
state_dict
,例如state_dict = net.state_dict()
。 - 使用
torch.save()
将state_dict
保存到文件。
- 实例化一个模型,例如
示例:
import torch
from torch import nn# 假设 MLP 已经被定义
net = MLP()
# 获取参数字典
state_dict = net.state_dict()
# 保存参数到文件
torch.save(state_dict, 'mlp-parameters.pt')


4.2.2 加载模型参数
load_state_dict()
:- 这是
nn.Module
的一个方法,它接收一个state_dict
字典,并用其内容覆盖当前模型的参数。
- 这是
- 加载步骤:
- 确保你有模型的代码定义。
- 实例化一个新的模型对象,例如
clone_net = MLP()
。此时,这个新模型的参数是随机初始化的。 - 使用
torch.load()
加载之前保存的state_dict
文件。 - 使用
clone_net.load_state_dict()
将加载的参数字典应用到新的模型实例上。
示例:
# 假设 MLP 的定义代码仍然可用
clone_net = MLP()
# 从文件加载参数字典
loaded_state_dict = torch.load('mlp-parameters.pt')
# 将参数加载到新模型中
clone_net.load_state_dict(loaded_state_dict)
验证:
- 通过比较加载前后模型的输出,可以验证参数是否成功加载。如果输入相同的随机数据,两个模型的输出应该完全一致。

5 使用 GPU
5.1 确认和查看 GPU 信息
在使用 GPU 之前,你需要确认系统是否拥有可用的 GPU,并查看其状态。
- 命令行工具:
- 在终端中运行
!nvidia-smi
(在 Jupyter 或 Colab 中)或直接运行nvidia-smi
。 - 这会显示你的 GPU 型号、内存使用情况、GPU 利用率(
GPU-Util
)以及正在运行的进程。 - GPU 内存(
GPU Memory
)和 CPU 内存是独立的。
- 在终端中运行
- PyTorch 方法:
torch.cuda.device_count()
: 返回可用的 GPU 数量。如果为 0,则表示没有可用的 GPU。

5.2 指定计算设备 (Device)
PyTorch 中的所有计算默认都在 CPU 上进行。要使用 GPU,你需要将数据和模型移动到 GPU 设备上。
-
定义设备:
torch.device('cpu')
: 指定使用 CPU。torch.device('cuda')
: 指定使用默认的 GPU(通常是第 0 号)。torch.device('cuda:1')
: 指定使用第 1 号 GPU。
-
实用函数:
- 一个好的实践是编写一个函数来自动选择可用的 GPU。
- 例如,
try_gpu(i=0)
尝试返回第i
号 GPU,如果不存在,则返回 CPU。
5.3 张量 (Tensor) 与 GPU
张量需要在 GPU 上才能利用 GPU 加速进行计算。
5.3.1 创建 GPU 上的张量
有两种方法可以将张量放到 GPU 上:
- 在创建时指定设备:
x = torch.tensor([0, 1, 2, 3], device=try_gpu(0))
- 这会直接在第 0 号 GPU 的内存上创建张量
x
。
- 将现有张量移动到 GPU:
y = torch.randn(2, 3)
y = y.to(try_gpu(1))
或y = y.cuda(1)
- 这会把张量
y
从 CPU 内存复制到第 1 号 GPU 内存。

5.3.2 GPU 上的张量运算
- 一致性原则: 所有参与运算的张量必须位于同一个设备上。
- 例如,如果
x
在第 0 号 GPU,y
在第 1 号 GPU,直接计算x + y
会报错。 - 你需要手动将它们移动到同一个 GPU 上,例如
y_on_gpu0 = y.to(try_gpu(0))
,然后计算x + y_on_gpu0
。
为什么必须手动移动?
这是为了性能考虑。在不同设备(特别是 CPU 和 GPU)之间传输数据是一项耗时的操作。PyTorch 强制你明确指定设备,可以帮助你避免因不经意的数据移动而造成的性能瓶颈。

5.4 神经网络与 GPU
将模型放在 GPU 上进行训练与处理张量类似,通过 .to()
方法实现。
- 步骤:
- 创建网络:
net = MLP()
,模型默认在 CPU 上。 - 移动网络到 GPU:
net.to(try_gpu(0))
。 - 移动数据到 GPU:
X = X.to(try_gpu(0))
。 - 进行前向计算:
net(X)
。
- 创建网络:
- 验证:
- 你可以通过检查模型参数的
.device
属性来确认模型是否在 GPU 上,例如net[0].weight.device
。 - 如果模型和数据都在同一个 GPU 上,前向和反向传播计算都会在该 GPU 上自动完成。
- 你可以通过检查模型参数的
