当前位置：首页 > news >正文

PyTorch：主要模块简介

news 2026/5/24 4:27:03

PyTorch 是 Python 生态中最常用的深度学习框架之一。它并不是只提供若干神经网络层，也不是简单的 GPU 计算工具，而是围绕深度学习的完整流程，组织出一组相互配合的模块：张量计算、自动求导、神经网络建模、数据加载、优化器、设备加速、模型保存、分布式训练与部署扩展等。

在学习 PyTorch 的过程中，如果一时记不清有哪些顶层名称，可以直接在 Python 中查看：

import torch print(dir(torch))

输出中通常可以看到许多顶层名称，例如：Tensor、nn、autograd、optim、utils、cuda、mps、amp、linalg、fft、sparse、distributed、jit、export、profiler 等。需要注意的是，dir(torch) 显示的是顶层名称集合，并不等于学习时必须逐一掌握的模块清单。

从整体上看，可以把 PyTorch 的主要模块理解为三类分工：

第一类是核心训练工作流模块，负责张量表示、自动求导、模型定义、数据加载、参数优化和设备加速。

第二类是模型构建与计算支持模块，负责神经网络层、函数式接口、线性代数、概率分布、稀疏张量、傅里叶变换等计算能力。

第三类是工程扩展与生态模块，负责编译加速、模型导出、分布式训练、性能分析、日志可视化和领域扩展。

这种模块化设计使得 PyTorch 不只是“神经网络层的集合”，而是一个围绕深度学习训练与部署组织起来的完整计算框架。

一、第一层：核心训练工作流模块

这一层模块直接对应深度学习最常见的训练流程：准备数据、定义模型、前向计算、计算损失、反向传播、更新参数、保存模型。在初学阶段，这一层最值得优先掌握。

1、torch：张量与基础计算模块

torch 是 PyTorch 的核心模块。它提供张量对象、数学运算、数组变形、索引切片、随机数生成、设备迁移等基础能力。可以把它理解为 PyTorch 中最底层、最常用的计算入口。

张量（Tensor）是 PyTorch 中最基本的数据结构。输入数据、模型参数、中间激活值、梯度、损失值，最终都以张量形式参与计算。

常用工具：

• torch.tensor()：从 Python 列表、数值等对象创建张量

• torch.zeros()：创建全 0 张量

• torch.ones()：创建全 1 张量

• torch.randn()：创建服从标准正态分布的随机张量

• torch.arange()：创建等差序列张量

• torch.cat()：沿指定维度拼接张量

• torch.stack()：把多个张量堆叠成新维度

• torch.reshape()：改变张量形状

• torch.matmul()：矩阵乘法

• torch.save()：保存张量、模型参数或其他对象

• torch.load()：加载已保存的对象

示例：

import torch # 输入矩阵 (2,2)x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])# 权重矩阵 (2,1)，随机初始化w = torch.randn(2, 1) # 矩阵乘法：y = x @ w，输出形状 (2,1)y = torch.matmul(x, w) print("x 的形状：", x.shape) # (2,2)print("w 的形状：", w.shape) # (2,1)print("y 的形状：", y.shape) # (2,1)

这个示例说明：在 PyTorch 中，很多深度学习计算本质上都是张量之间的数学运算。

2、torch.autograd：自动求导模块

torch.autograd 是 PyTorch 的自动求导系统。它负责记录张量运算过程，并根据计算图自动计算梯度。官方文档将 torch.autograd 描述为用于实现任意标量值函数自动微分的类与函数集合。

在深度学习中，模型训练依赖反向传播。手动推导和实现每个参数的梯度既复杂又容易出错，而 autograd 可以根据前向计算过程自动完成梯度计算。

常用工具：

• requires_grad=True：表示需要跟踪该张量的梯度

• loss.backward()：从损失值出发执行反向传播

• tensor.grad：查看张量对应的梯度

• torch.no_grad()：关闭梯度记录，常用于推理阶段

• tensor.detach()：从当前计算图中分离张量

示例：

import torch # 定义需要梯度的标量 x = 2.0x = torch.tensor(2.0, requires_grad=True) # 计算 y = x² + 3x + 1y = x ** 2 + 3 * x + 1 # 反向传播，自动计算 y 对 x 的导数y.backward() print("y 的值：", y.item()) # 输出 y 的计算结果print("dy/dx：", x.grad.item()) # 输出导数 (2x + 3) 在 x=2 时的值

这个示例中，PyTorch 会自动计算 y 关于 x 的导数。对于神经网络而言，模型参数通常都会参与自动求导，训练时不需要手动写出每个参数的梯度公式。

3、torch.nn：神经网络模块

torch.nn 是 PyTorch 中构建神经网络最核心的模块。它提供模型基类、常用网络层、激活函数、损失函数、容器结构等内容。官方文档中，torch.nn.Module 是所有神经网络模块的基类；PyTorch 也使用模块来表示神经网络。

在 PyTorch 中，通常会通过继承 nn.Module 来定义模型。一个模型一般包含两部分：

• 在 __init__() 中定义网络层

• 在 forward() 中定义前向传播逻辑

常用工具：

• nn.Module：所有神经网络模型的基类

• nn.Linear：全连接层

• nn.Conv2d：二维卷积层

• nn.MaxPool2d：二维最大池化层

• nn.BatchNorm2d：二维批量归一化层

• nn.LayerNorm：层归一化

• nn.Dropout：随机失活层

• nn.Embedding：词嵌入层

• nn.RNN：循环神经网络层

• nn.LSTM：长短期记忆网络层

• nn.GRU：门控循环单元

• nn.Transformer：Transformer 结构模块

• nn.CrossEntropyLoss：交叉熵损失

• nn.MSELoss：均方误差损失

示例：

import torch # PyTorch 框架import torch.nn as nn # 神经网络模块 # 定义简单前馈网络class SimpleNet(nn.Module): def __init__(self): super().__init__() # 顺序容器：全连接 → ReLU → 全连接 self.net = nn.Sequential( nn.Linear(4, 8), # 输入4维，输出8维 nn.ReLU(), # ReLU激活 nn.Linear(8, 3) # 输出3维（分类logits） ) def forward(self, x): return self.net(x) # 前向传播 model = SimpleNet() # 实例化模型 # 批量输入：5个样本，每个4个特征x = torch.randn(5, 4) output = model(x) # 前向传播得到输出 print("输出形状：", output.shape) # torch.Size([5, 3])

这个示例体现了 PyTorch 建模的基本思想：模型是一个 nn.Module，输入张量经过 forward() 得到输出张量。

4、torch.nn.functional：函数式神经网络接口

torch.nn.functional 通常简写为 F。它提供许多函数式操作，例如激活函数、池化、归一化、损失函数等。

torch.nn 与 torch.nn.functional 的区别可以简单理解为：

• torch.nn 更偏“层对象”，适合定义带参数、需要注册到模型中的模块

• torch.nn.functional 更偏“函数调用”，适合没有参数或临时调用的计算操作

常用工具：

• F.relu()：ReLU 激活函数

• F.softmax()：Softmax 函数

• F.log_softmax()：对数 Softmax

• F.cross_entropy()：交叉熵损失

• F.mse_loss()：均方误差损失

• F.dropout()：函数式随机失活

• F.conv2d()：函数式二维卷积

• F.max_pool2d()：函数式二维最大池化

示例：

import torchimport torch.nn as nnimport torch.nn.functional as F class SimpleClassifier(nn.Module): """一个简单的全连接分类器：输入4维，输出3维（类别logits）""" def __init__(self): super().__init__() # 第一层全连接：输入4 → 输出8 self.fc1 = nn.Linear(4, 8) # 第二层全连接：输入8 → 输出3（类别数） self.fc2 = nn.Linear(8, 3) def forward(self, x): # 第一层输出经过ReLU激活 x = F.relu(self.fc1(x)) # 第二层输出（logits，无激活） x = self.fc2(x) return x # 实例化模型model = SimpleClassifier()

这个示例中，nn.Linear 用于定义带参数的线性层，而 F.relu() 用于执行无参数的激活操作。

5、torch.optim：优化器模块

torch.optim 用于实现参数优化算法。神经网络训练的目标是让损失函数逐步减小，而优化器负责根据梯度更新模型参数。官方文档将 torch.optim 描述为实现各种优化算法的包，并支持常见优化方法。

常用工具：

• optim.SGD：随机梯度下降

• optim.Adam：自适应矩估计优化器，深度学习中非常常用

• optim.AdamW：解耦权重衰减版本的 Adam，训练 Transformer 类模型时很常见

• optim.RMSprop：基于平方梯度移动平均的优化器

• optim.Adadelta：自适应学习率优化器

• optim.lr_scheduler：学习率调度器子模块

训练时常见的优化步骤是：

optimizer.zero_grad()loss.backward()optimizer.step()

含义分别是：

• 清空上一轮梯度

• 根据损失执行反向传播

• 按照优化算法更新参数

示例：

import torchimport torch.nn as nnimport torch.optim as optim # 线性回归模型：输入1维，输出1维model = nn.Linear(1, 1)# 均方误差损失函数criterion = nn.MSELoss()# SGD优化器，学习率0.01optimizer = optim.SGD(model.parameters(), lr=0.01) # 随机生成20个输入样本（批量）x = torch.randn(20, 1)# 目标值：y = 3x + 2y = 3 * x + 2 # 训练100轮for epoch in range(100): pred = model(x) # 前向传播：计算预测值 loss = criterion(pred, y) # 计算损失 optimizer.zero_grad() # 清空梯度 loss.backward() # 反向传播 optimizer.step() # 更新参数 print("权重：", model.weight.item()) # 应接近3print("偏置：", model.bias.item()) # 应接近2

这个示例展示了 PyTorch 最小训练闭环：

前向计算 → 损失计算 → 反向传播 → 参数更新

6、torch.utils.data：数据加载模块

torch.utils.data 主要用于组织数据集和批量加载数据。深度学习训练通常不是一次性把全部数据送入模型，而是按批次读取、打乱、并行加载。官方文档指出，DataLoader 是 PyTorch 数据加载工具的核心，它表示一个基于数据集的 Python 可迭代对象。

常用工具：

• Dataset：自定义数据集的基类

• TensorDataset：把张量包装成数据集

• DataLoader：按批次加载数据

• random_split()：随机划分数据集

• Sampler：控制样本抽样方式

• WeightedRandomSampler：按权重抽样，常用于类别不平衡任务

示例：

import torchfrom torch.utils.data import TensorDataset, DataLoader # 随机生成100个样本，每个样本4个特征X = torch.randn(100, 4)# 随机生成100个标签（0、1、2三类）y = torch.randint(0, 3, (100,)) # 封装为数据集：特征张量和标签张量dataset = TensorDataset(X, y)# 创建DataLoader：批量大小16，打乱顺序loader = DataLoader(dataset, batch_size=16, shuffle=True) # 遍历一个batch，查看形状并退出（仅演示）for batch_X, batch_y in loader: print(batch_X.shape, batch_y.shape) # (16,4) (16,) break

这个示例说明：Dataset 负责表示“有哪些样本”，DataLoader 负责表示“如何按批次读取样本”。

7、torch.cuda、torch.mps 与 torch.amp：设备加速模块

深度学习训练通常依赖硬件加速。PyTorch 可以把张量和模型移动到不同设备上，例如 CPU、NVIDIA GPU、Apple Silicon GPU 等。

常用工具：

• torch.cuda.is_available()：检查 CUDA 是否可用

• torch.device()：构建设备对象

• tensor.to(device)：把张量移动到指定设备

• model.to(device)：把模型参数移动到指定设备

• torch.cuda.empty_cache()：释放 CUDA 缓存

• torch.mps：用于访问 Apple Metal Performance Shaders 后端，适合部分 macOS 设备

• torch.amp：自动混合精度模块，可让部分操作使用低精度类型，以提升训练速度并降低显存占用

示例：

import torchimport torch.nn as nn # 检测可用设备：优先CUDA（NVIDIA GPU），其次MPS（Apple Silicon GPU），最后CPUif torch.cuda.is_available(): device = torch.device("cuda")elif torch.backends.mps.is_available(): device = torch.device("mps")else: device = torch.device("cpu") # 将模型和数据移动到选定的设备上model = nn.Linear(4, 3).to(device)x = torch.randn(8, 4).to(device) output = model(x) # 前向传播在设备上执行 print("当前设备：", device)print("输出设备：", output.device) # 确认输出张量所在的设备

这个示例体现了 PyTorch 设备管理的基本方式：模型和数据必须放在同一设备上，才能正常参与计算。

综合示例 1：从数据到模型训练的最小闭环

下面这个示例把第一层中的几个核心模块串起来：先构造数据，再加载批次，然后定义模型、损失函数和优化器，最后完成简单训练。

import torchimport torch.nn as nnimport torch.optim as optimfrom torch.utils.data import TensorDataset, DataLoader # 1. 构造示例数据X = torch.randn(200, 4) # 200个样本，每个4个特征y = torch.randint(0, 3, (200,)) # 随机标签（3类） dataset = TensorDataset(X, y) # 封装为数据集loader = DataLoader(dataset, batch_size=32, shuffle=True) # 批量加载，打乱顺序 # 2. 定义模型model = nn.Sequential( nn.Linear(4, 16), # 输入4维，输出16维 nn.ReLU(), # ReLU激活 nn.Linear(16, 3) # 输出3维（类别logits）) # 3. 定义损失函数与优化器criterion = nn.CrossEntropyLoss() # 交叉熵损失（多分类）optimizer = optim.Adam(model.parameters(), lr=0.01) # Adam优化器 # 4. 训练模型for epoch in range(5): # 迭代5个epoch total_loss = 0.0 for batch_X, batch_y in loader: pred = model(batch_X) # 前向传播 loss = criterion(pred, batch_y) # 计算损失 optimizer.zero_grad() # 清空梯度 loss.backward() # 反向传播 optimizer.step() # 更新参数 total_loss += loss.item() # 累加当前batch损失 print(f"第 {epoch + 1} 轮，损失：{total_loss:.4f}")

这个示例体现了 PyTorch 工作流的典型顺序：

数据准备 → 批量加载 → 模型定义 → 前向计算 → 损失计算 → 反向传播 → 参数更新

二、第二层：模型构建与计算支持模块

这一层模块主要负责“模型内部如何计算”。如果说第一层解决的是“怎样完成训练流程”，那么第二层解决的是“模型需要哪些计算能力”。

1、torch.nn 中的层、容器与损失函数

torch.nn 不只是提供 nn.Module 基类，还提供了大量常用网络层、容器和损失函数。

常用层：

• nn.Linear：全连接层，常用于 MLP、分类头和回归头

• nn.Conv1d：一维卷积，常用于序列信号

• nn.Conv2d：二维卷积，常用于图像任务

• nn.Conv3d：三维卷积，常用于视频或体数据

• nn.BatchNorm1d、nn.BatchNorm2d：批量归一化

• nn.LayerNorm：层归一化，Transformer 中常见

• nn.Dropout：随机失活，常用于缓解过拟合

• nn.Embedding：嵌入层，常用于文本、推荐系统等任务

• nn.MultiheadAttention：多头注意力层

常用容器：

• nn.Sequential：按顺序组织多个层

• nn.ModuleList：保存多个子模块

• nn.ModuleDict：用字典形式保存多个子模块

常用损失函数：

• nn.CrossEntropyLoss：多分类交叉熵损失

• nn.BCELoss：二分类交叉熵损失

• nn.BCEWithLogitsLoss：带 Logits 的二分类交叉熵损失

• nn.MSELoss：均方误差损失

• nn.L1Loss：平均绝对误差损失

• nn.NLLLoss：负对数似然损失

2、torch.linalg：线性代数模块

torch.linalg 用于线性代数计算。它在机器学习、深度学习、科学计算中都非常重要，尤其适合矩阵分解、范数、矩阵秩、逆矩阵、特征值等计算。官方文档中，torch.linalg 包含矩阵秩、条件数等典型线性代数函数。

常用工具：

• torch.linalg.norm()：计算向量或矩阵范数

• torch.linalg.inv()：计算矩阵逆

• torch.linalg.det()：计算行列式

• torch.linalg.matrix_rank()：计算矩阵秩

• torch.linalg.eig()：计算特征值和特征向量

• torch.linalg.svd()：奇异值分解

• torch.linalg.qr()：QR 分解

示例：

import torch # 创建一个随机 3x3 矩阵A = torch.randn(3, 3) # 计算奇异值分解 (SVD)，返回 U, 奇异值向量, V^Hu, s, vh = torch.linalg.svd(A) # 打印奇异值（降序排列）print("奇异值：", s)

这一模块更偏数学计算基础。在实现 PCA、矩阵分解、低秩近似等算法时，它非常有用。

3、torch.fft：傅里叶变换模块

torch.fft 用于离散傅里叶变换及相关计算。官方文档将 torch.fft 定位为离散傅里叶变换和相关函数集合。

常用工具：

• torch.fft.fft()：一维快速傅里叶变换

• torch.fft.ifft()：一维逆傅里叶变换

• torch.fft.fft2()：二维快速傅里叶变换

• torch.fft.ifft2()：二维逆傅里叶变换

• torch.fft.rfft()：实值输入的一维傅里叶变换

• torch.fft.fftfreq()：生成傅里叶频率

应用场景：

• 信号处理

• 图像频域分析

• 科学计算

• 某些频域神经网络结构

4、torch.sparse：稀疏张量模块

torch.sparse 用于处理稀疏张量。稀疏张量适合表示大多数元素为 0 的数据，例如图结构、推荐系统中的用户—物品矩阵、高维文本特征等。官方文档中，torch.sparse 提供从稠密张量构造稀疏张量以及多种稀疏布局支持。

常用工具：

• torch.sparse_coo_tensor()：创建 COO 格式稀疏张量

• torch.sparse_csr_tensor()：创建 CSR 格式稀疏张量

• tensor.to_sparse()：把稠密张量转换为稀疏张量

• tensor.to_dense()：把稀疏张量转换为稠密张量

• torch.sparse.mm()：稀疏矩阵乘法

• torch.sparse.softmax()：稀疏版本 Softmax

示例：

import torch # 显式启用稀疏张量不变性检查（避免警告，确保格式正确）with torch.sparse.check_sparse_tensor_invariants(True): # 定义非零元素的坐标（每列为一个坐标对，共3个非零元） indices = torch.tensor([[0, 1, 1], [2, 0, 2]]) # 对应坐标上的数值 values = torch.tensor([3.0, 4.0, 5.0]) # 创建 2x3 的 COO 格式稀疏张量 x = torch.sparse_coo_tensor(indices, values, size=(2, 3)) # 打印稀疏表示（坐标和值）print(x)# 转换为稠密张量并打印print(x.to_dense())

稀疏张量并不是初学阶段的重点，但在图神经网络、推荐系统和大规模特征处理中很重要。

5、torch.distributions：概率分布模块

torch.distributions 用于表示概率分布、采样和计算概率值。它常用于生成模型、强化学习、贝叶斯建模和变分推断等任务。

常用工具：

• Normal：正态分布

• Bernoulli：伯努利分布

• Categorical：类别分布

• MultivariateNormal：多元正态分布

• Uniform：均匀分布

• Distribution.sample()：从分布中采样

• Distribution.log_prob()：计算对数概率

示例：

import torchfrom torch.distributions import Normal # 创建一个标准正态分布（均值为0，标准差为1）dist = Normal(loc=0.0, scale=1.0) # 从分布中采样5个样本x = dist.sample((5,))# 计算每个样本的对数概率密度（log PDF）log_p = dist.log_prob(x) print("采样结果：", x)print("对数概率：", log_p)

在普通监督学习中，可能很少直接使用这个模块；但在 VAE、策略梯度、概率建模中，它非常常见。

6、torch.random：随机数模块

torch.random 与 PyTorch 的随机数状态有关。深度学习实验通常需要控制随机性，例如初始化参数、打乱数据、随机增强、Dropout 等。

常用工具：

• torch.manual_seed()：设置随机种子

• torch.initial_seed()：查看初始随机种子

• torch.get_rng_state()：获取随机数生成器状态

• torch.set_rng_state()：设置随机数生成器状态

示例：

import torch # 固定随机种子，使得后续随机操作可复现torch.manual_seed(42) # 生成一个长度为3的随机张量（标准正态分布）x1 = torch.randn(3) # 再次设置相同的随机种子，确保接下来生成的随机数与 x1 相同torch.manual_seed(42) x2 = torch.randn(3) # 与 x1 相同的随机数序列 print(x1)print(x2) # 输出应与 x1 完全一致

这个示例说明：设置相同随机种子后，相同随机操作可以得到一致结果。需要注意的是，完整可复现实验还可能涉及 CUDA、DataLoader 多进程和底层算子确定性等因素。

7、torch.func：函数变换模块

torch.func 提供函数式变换能力，例如自动向量化、函数式梯度、雅可比矩阵、海森矩阵等。它更适合进阶场景，尤其是元学习、高阶梯度、批量函数变换和研究型代码。

常用工具：

• torch.func.grad()：构造梯度函数

• torch.func.vmap()：自动向量化映射

• torch.func.jacrev()：反向模式计算雅可比矩阵

• torch.func.jacfwd()：前向模式计算雅可比矩阵

• torch.func.hessian()：计算海森矩阵

初学阶段只需要知道：torch.func 不是普通训练循环的必需模块，而是面向更高级自动微分与函数变换的工具。

综合示例 2：定义一个更标准的训练函数

下面的示例展示 PyTorch 中较常见的训练组织方式：把单轮训练封装为函数，使代码更接近真实项目结构。

import torchimport torch.nn as nnimport torch.optim as optimfrom torch.utils.data import TensorDataset, DataLoader # 定义一个多层感知机（MLP）分类器class MLP(nn.Module): def __init__(self, input_dim=4, hidden_dim=16, output_dim=3): super().__init__() # 使用Sequential构建网络：输入层 → 隐藏层（ReLU）→ 输出层 self.net = nn.Sequential( nn.Linear(input_dim, hidden_dim), # 全连接层：4 → 16 nn.ReLU(), # ReLU激活 nn.Linear(hidden_dim, output_dim) # 输出层：16 → 3（logits） ) def forward(self, x): return self.net(x) # 前向传播 # 定义单个epoch的训练函数def train_one_epoch(model, loader, criterion, optimizer): model.train() # 设置为训练模式 total_loss = 0.0 for batch_X, batch_y in loader: pred = model(batch_X) # 前向传播得到预测logits loss = criterion(pred, batch_y) # 计算损失 optimizer.zero_grad() # 清空旧梯度 loss.backward() # 反向传播 optimizer.step() # 更新参数 total_loss += loss.item() # 累加当前batch损失 return total_loss / len(loader) # 返回平均损失 # 构造数据：200个样本，每个4个特征，3类标签（0,1,2）X = torch.randn(200, 4)y = torch.randint(0, 3, (200,)) # 创建DataLoader：批量32，打乱顺序loader = DataLoader( TensorDataset(X, y), batch_size=32, shuffle=True) # 实例化模型、损失函数、优化器model = MLP()criterion = nn.CrossEntropyLoss() # 交叉熵损失（多分类）optimizer = optim.Adam(model.parameters(), lr=0.01) # Adam优化器 # 训练5个epochfor epoch in range(5): avg_loss = train_one_epoch(model, loader, criterion, optimizer) print(f"第 {epoch + 1} 轮，平均损失：{avg_loss:.4f}")

这段代码体现了 PyTorch 训练代码的基本组织方式：模型类、训练函数、数据加载器、损失函数和优化器彼此分工明确。

三、第三层：工程扩展与生态模块

这一层模块同样重要，但专题性更强。初学阶段不必全部展开，只需要知道它们分别解决什么问题。等到模型规模变大、训练变慢、需要部署或需要多卡训练时，再逐步深入。

1、torch.compile 与 torch.compiler：编译加速模块

torch.compile() 是 PyTorch 2.x 之后非常重要的加速入口。它可以尝试把模型计算图编译成更高效的执行形式，从而提升训练或推理速度。torch.compiler 则提供与编译相关的 API。

常用方式：

model = torch.compile(model)

适用场景：

• 模型训练速度优化

• 推理性能优化

• 在不大幅改变原有代码的情况下尝试加速

需要注意的是，编译加速并不保证所有模型都一定更快，也可能受到模型结构、硬件设备、算子支持和输入形状变化的影响。

2、torch.export：模型导出模块

torch.export 用于捕获模型的张量计算图，并生成可进一步处理的导出表示。官方文档中，torch.export.export() 会基于示例输入追踪模型执行，并记录 PyTorch 操作；相比传统追踪方式，它更强调对动态形状和程序语义的处理能力。

常用场景：

• 模型部署前的图捕获

• 编译器后端处理

• 模型转换和进一步优化

• 与 PyTorch 2.x 编译生态配合使用

初学阶段可以先知道：torch.export 更偏工程部署和编译链路，不是普通训练入门的第一重点。

3、torch.jit：TorchScript 模块

torch.jit 是较早期 PyTorch 中用于模型脚本化和图模式表示的重要模块。它可以把部分 PyTorch 模型转换为 TorchScript 形式，以便脱离 Python 环境执行或做进一步优化。

常用工具：

• torch.jit.trace()：通过示例输入追踪模型执行路径

• torch.jit.script()：把支持的 Python 模型代码脚本化

• torch.jit.save()：保存 TorchScript 模型

• torch.jit.load()：加载 TorchScript 模型

在当前 PyTorch 生态中，torch.compile、torch.export 等工具越来越重要，但理解 torch.jit 仍有助于理解 PyTorch 从动态图到图表示的发展脉络。

4、torch.distributed：分布式训练模块

torch.distributed 用于多进程、多 GPU、多节点训练。官方文档指出，该包为跨多个计算节点运行的多进程并行提供 PyTorch 支持和通信原语。

常用工具：

• torch.distributed.init_process_group()：初始化分布式进程组

• torch.distributed.get_rank()：获取当前进程编号

• torch.distributed.get_world_size()：获取总进程数

• DistributedDataParallel：分布式数据并行训练，通常简写为 DDP

• DistributedSampler：分布式数据采样器

典型应用场景：

• 单机多卡训练

• 多机多卡训练

• 大模型训练

• 大规模图像、语音、文本模型训练

对于初学者而言，先掌握单机单卡训练即可；当模型和数据规模变大时，再学习 DDP 会更自然。

5、torch.profiler：性能分析模块

torch.profiler 用于分析模型训练或推理过程中的性能瓶颈。官方文档中，torch.profiler 可导出设备上的内存事件信息和时间线图等内容。

常用用途：

• 分析 CPU 和 GPU 时间消耗

• 观察哪些算子最耗时

• 分析显存占用

• 配合 TensorBoard 查看性能轨迹

常用工具：

• torch.profiler.profile()：启动性能分析

• torch.profiler.record_function()：为代码块添加分析标签

• ProfilerActivity.CPU：记录 CPU 活动

• ProfilerActivity.CUDA：记录 CUDA 活动

性能分析通常不是写第一个模型时需要掌握的内容，但在模型训练变慢、显存占用异常或推理延迟过高时非常重要。

6、torch.utils.tensorboard：日志与可视化模块

torch.utils.tensorboard 用于把训练过程中的指标写入 TensorBoard。例如损失曲线、准确率曲线、学习率变化、图像样本、模型结构等。

常用工具：

• SummaryWriter：写入 TensorBoard 日志

• add_scalar()：记录标量指标

• add_image()：记录图像

• add_histogram()：记录参数或梯度分布

• add_graph()：记录模型图

示例：

from torch.utils.tensorboard import SummaryWriter # 创建TensorBoard写入器，日志将保存在 runs/demo 目录下writer = SummaryWriter("runs/demo") # 模拟训练过程，记录每个epoch的损失值for epoch in range(10): # 模拟损失：随epoch增加而减小 loss = 1.0 / (epoch + 1) # 记录标量数据：标签为"Loss/train"，值为loss，步数为epoch writer.add_scalar("Loss/train", loss, epoch) # 关闭写入器，释放资源writer.close()

这个模块适合在训练过程中观察模型是否正常收敛，而不是只在训练结束后看最终结果。

7、torch.package 与 torch.hub：模型打包与模型获取模块

torch.package 用于创建包含 PyTorch 代码和模型工件的包，便于保存、共享和加载。官方文档将其描述为支持创建包含模型制品和任意 PyTorch 代码的包。

torch.hub 则更偏向从远程仓库加载模型或代码。早期很多预训练模型会通过 torch.hub 提供加载入口。

常用场景：

• 保存和共享模型相关代码

• 加载已有模型

• 快速复现实验

• 组织模型发布流程

需要注意的是，实际生产部署中，模型格式、依赖环境、推理框架和安全策略都需要进一步考虑，不能只依赖单一保存函数。

8、torchvision、torchaudio 与其他生态库

PyTorch 本体之外，还有一组常见生态库，用于处理不同领域的数据和模型。PyTorch 官网也强调其生态工具和库支持计算机视觉、自然语言处理等方向，并支持分布式训练和生产部署。

常见生态库：

• torchvision：计算机视觉库，提供图像数据集、图像变换、视觉模型等

• torchaudio：音频处理库，提供音频数据加载、特征提取和音频模型相关工具

• torchtext：文本处理相关工具，但当前使用时需要注意版本兼容和维护状态

• torchmetrics：常用评价指标库，常与 PyTorch Lightning 配合使用

• torchao：面向量化、低精度和推理优化的生态工具

• TorchRL：强化学习相关工具库

这些库不一定都属于 torch 顶层模块，但在真实项目中经常与 PyTorch 一起使用。

综合示例 3：保存与加载模型参数

下面展示 PyTorch 中最常见的模型保存方式：保存 state_dict，即模型参数字典。

import torchimport torch.nn as nn # 定义一个简单模型：输入4维 → 16维 → ReLU → 输出3维model = nn.Sequential( nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3)) # 保存模型的参数（state_dict）到文件 model.pthtorch.save(model.state_dict(), "model.pth") # 重新创建一个结构完全相同的模型new_model = nn.Sequential( nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3)) # 从文件加载参数到新模型state_dict = torch.load("model.pth")new_model.load_state_dict(state_dict) # 切换到推理模式（关闭Dropout/BatchNorm等训练行为）new_model.eval() print("模型参数已加载")

在 PyTorch 中，推荐把“模型结构定义代码”和“模型参数 state_dict”配合使用。这样更清晰，也更符合常见工程实践。

四、学习 PyTorch 模块时的优先顺序

PyTorch 模块很多，但初学时不必平均用力。更合理的学习顺序是先掌握训练闭环，再逐步扩展到复杂模型和工程能力。

第一阶段，应优先掌握：

• torch：张量创建、张量运算、形状变换、设备迁移。

• torch.autograd：自动求导、backward()、梯度概念。

• torch.nn：模型定义、常用层、损失函数。

• torch.optim：优化器、参数更新。

• torch.utils.data：数据集与数据加载器。

第二阶段，可以继续掌握：

• torch.nn.functional：函数式接口。

• torch.cuda、torch.mps、torch.amp：设备与混合精度训练。