PyTorch：主要模块简介

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_gradTrue表示需要跟踪该张量的梯度• loss.backward()从损失值出发执行反向传播• tensor.grad查看张量对应的梯度• torch.no_grad()关闭梯度记录常用于推理阶段• tensor.detach()从当前计算图中分离张量示例import torch # 定义需要梯度的标量 x 2.0x torch.tensor(2.0, requires_gradTrue) # 计算 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) 在 x2 时的值这个示例中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.TransformerTransformer 结构模块• 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(), lr0.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_size16, shuffleTrue) # 遍历一个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 # 检测可用设备优先CUDANVIDIA GPU其次MPSApple 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_size32, shuffleTrue) # 批量加载打乱顺序 # 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(), lr0.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标准差为1dist Normal(loc0.0, scale1.0) # 从分布中采样5个样本x dist.sample((5,))# 计算每个样本的对数概率密度log PDFlog_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_dim4, hidden_dim16, output_dim3): 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 → 3logits ) 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,2X torch.randn(200, 4)y torch.randint(0, 3, (200,)) # 创建DataLoader批量32打乱顺序loader DataLoader( TensorDataset(X, y), batch_size32, shuffleTrue) # 实例化模型、损失函数、优化器model MLP()criterion nn.CrossEntropyLoss() # 交叉熵损失多分类optimizer optim.Adam(model.parameters(), lr0.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.jitTorchScript 模块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设备与混合精度训练。• 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 的结构并逐步建立完整的深度学习训练与工程实践意识。“点赞有美意赞赏是鼓励”