PyTorch：主要模块简介

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

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

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()：加载已保存的对象

示例：

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

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

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

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

常用工具：

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

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

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

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

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

示例：

这个示例中，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：均方误差损失

示例：

这个示例体现了 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()：函数式二维最大池化

示例：

这个示例中，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.step()

含义分别是：

• 清空上一轮梯度

• 根据损失执行反向传播

• 按照优化算法更新参数

示例：

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

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

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

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

常用工具：

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

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

• DataLoader：按批次加载数据

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

• Sampler：控制样本抽样方式

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

示例：

这个示例说明：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：自动混合精度模块，可让部分操作使用低精度类型，以提升训练速度并降低显存占用

示例：

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

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

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

这个示例体现了 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 分解

示例：

这一模块更偏数学计算基础。在实现 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

示例：

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

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

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

常用工具：

• Normal：正态分布

• Bernoulli：伯努利分布

• Categorical：类别分布

• MultivariateNormal：多元正态分布

• Uniform：均匀分布

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

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

示例：

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

6、torch.random：随机数模块

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

常用工具：

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

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

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

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

示例：

这个示例说明：设置相同随机种子后，相同随机操作可以得到一致结果。需要注意的是，完整可复现实验还可能涉及 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 中较常见的训练组织方式：把单轮训练封装为函数，使代码更接近真实项目结构。

这段代码体现了 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()：记录模型图

示例：

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

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，即模型参数字典。

在 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：设备与混合精度训练。

• torch.linalg：矩阵计算。

• torch.distributions：概率建模。

• torch.utils.tensorboard：训练日志记录。

第三阶段，再根据实际需要学习：

• torch.compile：模型编译加速。

• torch.export：模型导出。

• torch.distributed：分布式训练。

• torch.profiler：性能分析。

• torch.sparse：稀疏计算。

• torchvision、torchaudio 等领域生态库。

这种顺序符合深度学习项目的自然发展路径：先能训练一个模型，再能组织复杂模型，最后再考虑加速、部署和规模化训练。

五、PyTorch 与 Scikit-learn 模块组织方式的差异

Scikit-learn 更强调传统机器学习中的统一接口。许多模型都遵循 fit()、predict()、transform() 等接口，因此可以方便地替换模型、组合流程和比较结果。

PyTorch 的重点不同。它更强调动态计算图、张量计算、自动求导和灵活的模型定义方式。PyTorch 不会把完整训练流程封装成一个统一的 fit() 方法，而是把更多控制权交给开发者：用户自己写前向传播、损失计算、反向传播、参数更新和评估逻辑。

因此，二者的模块化思想有所不同：

• Scikit-learn 更像“规范化的机器学习工具箱”

• PyTorch 更像“可编程的深度学习计算框架”

• Scikit-learn 适合快速构建传统机器学习流程

• PyTorch 适合灵活构建神经网络、复杂训练过程和深度学习模型

这并不意味着二者互相替代。在真实项目中，Scikit-learn 常用于特征工程、传统机器学习基线和评估流程；PyTorch 则常用于深度神经网络、计算机视觉、自然语言处理、多模态模型和大规模训练。

小结

PyTorch 的主要模块大致可分为三层：核心训练工作流模块、模型构建与计算支持模块、工程扩展与生态模块。理解这些模块的分工，有助于从整体上把握 PyTorch 的结构，并逐步建立完整的深度学习训练与工程实践意识。

“点赞有美意，赞赏是鼓励”