PyTorch 进阶：数据集加载、模型训练与 GPU 加速
22.1 本章导学
上一章我们掌握了张量基础、自动微分与模型构建的核心语法，已经能够搭建简单的网络结构并完成单步训练。但真实的深度学习项目远不止单步迭代，它包含数据加载、多轮训练、验证评估、权重保存、GPU 加速、异常处理等完整的工程流程。一套规范的训练工程体系，不仅能提升训练效率，更能保证实验的可复现性与结果的可靠性，是从 “跑通 Demo” 到 “落地项目” 的关键跨越。
本章聚焦 PyTorch 工程化训练的全流程，沿着 “数据供给 - 训练循环 - 硬件加速 - 权重管理 - 优化技巧” 的脉络展开。首先讲解 Dataset 与 DataLoader 数据加载体系，掌握自定义数据集、批量加载、多进程加速的标准方法，这是大规模数据训练的基础；然后拆解完整的训练 - 验证闭环，建立规范的训练循环范式，掌握早停、指标记录等工程技巧；接着讲解 GPU 加速的核心方法，实现张量与模型的设备迁移，掌握基础的显存优化思路；最后覆盖模型保存加载、断点续训、学习率调度、梯度裁剪等工业级训练技巧。
所有内容均对应大模型开发的工程逻辑：大模型的分布式数据加载、训练循环、混合精度训练、权重保存，全部基于本章的基础范式扩展而来。吃透本章的标准训练流程，后续学习大模型微调、分布式训练时，就能快速理解高阶特性的底层逻辑。
22.2 数据集加载体系：Dataset 与 DataLoader
深度学习训练的第一步是数据供给。如果把所有数据一次性加载到内存，面对大规模数据集时会直接内存溢出；如果手动逐批次读取，不仅代码繁琐，还无法充分利用硬件性能。PyTorch 提供了标准化的数据加载接口，将数据集定义、采样、批量加载解耦，兼顾灵活性与效率。
22.2.1 Dataset：数据集的抽象基类
torch.utils.data.Dataset是所有数据集的抽象基类，它定义了数据访问的统一接口。自定义数据集只需要继承这个类，实现两个核心方法：
__len__：返回数据集的总样本数，让框架知道数据集的大小。__getitem__：接收一个索引，返回对应下标的单条样本，通常包含输入特征和标签。
这种设计的优势在于解耦了数据存储和训练流程。数据集可以存在硬盘、数据库或者网络上，只要通过__getitem__能读取到单条样本即可，不需要全部加载到内存中。大模型的海量语料训练，正是基于这种流式读取的思想，按需加载样本，避免内存溢出。
以文本分类数据集为例，自定义数据集的标准结构非常清晰：初始化时加载文件路径列表，获取样本时读取对应文件并做预处理，返回张量格式的特征和标签。这种写法内存占用极低，只在用到样本时才加载，适配任意规模的数据集。
22.2.2 DataLoader：批量加载与调度器
Dataset 只能读取单条样本，批量组织、打乱顺序、多进程加速这些工作，都由DataLoader完成。它接收一个 Dataset 实例，通过配置参数自动完成数据的批量加载。 最核心的几个参数直接决定了加载效率与训练效果：
batch_size：批次大小，每次迭代返回的样本数量。批次大小影响训练稳定性和显存占用，是训练的核心超参数之一。批次越大梯度越稳定，但显存占用越高，需要根据硬件显存调整。shuffle：是否在每轮训练前打乱数据集顺序。训练集通常设为 True，打乱样本顺序避免模型记住样本顺序，提升泛化能力；验证集和测试集设为 False，保证结果可复现。num_workers：加载数据的子进程数量。设为 0 表示只用主进程加载，数值越大加载速度越快，但会占用更多内存和 CPU 资源。合理设置 num_workers 可以让数据加载和模型训练并行，避免 GPU 等待数据，大幅提升硬件利用率。drop_last：是否丢弃最后一个不足一个批次的样本。当数据集大小不能被批次大小整除时，最后一批样本数较少，可能影响 BatchNorm 等层的统计量，训练时通常设为 True 丢弃。pin_memory：锁页内存，设为 True 可以加快张量从 CPU 迁移到 GPU 的速度，配合 GPU 训练时推荐开启。
22.2.3 自定义批处理：collate_fn
默认的 DataLoader 会直接把同批次的样本堆叠成张量，但很多场景下样本长度不一致，比如文本序列长度不同、图像尺寸不同，直接堆叠会报错。这时候就需要自定义collate_fn函数，在生成批次时做统一处理，比如对文本做填充对齐、对图像做缩放裁剪。collate_fn接收一个批次的样本列表，返回处理好的批量张量。在 NLP 任务中，它负责做序列填充、生成注意力掩码；在目标检测任务中，它负责对齐不同尺寸的标注。这是数据加载中最灵活的部分，也是处理变长数据的核心机制。大模型训练中的变长序列处理，底层就是通过自定义 collate_fn 实现的。
22.3 标准训练循环：训练与验证的完整闭环
22.3.1 为什么要分训练集和验证集
模型训练的目标是泛化能力，而不是在训练集上刷分。每轮训练结束后，在独立的验证集上评估效果，才能真实反映模型的泛化水平。验证集有三个核心作用：监控过拟合，当训练损失持续下降而验证损失上升时，说明已经过拟合；调优超参数，根据验证集效果调整学习率、批次大小等参数；早停机制的判断依据，验证效果不再提升时提前终止训练。 测试集则只用于最终评估，全程不能参与训练和调参，否则会出现数据泄露，评估结果虚高。
22.3.2 单轮训练的标准流程
每一轮训练都遵循固定的五步流程，这是所有深度学习训练的通用范式： 第一步，梯度清零。调用优化器的zero_grad()方法，清空上一轮迭代累积的梯度。PyTorch 默认梯度累加，如果不清零，梯度会不断叠加，导致更新方向错误。 第二步，前向传播。将批次数据输入模型，得到预测结果。这一步模型处于训练模式，Dropout、BatchNorm 等层正常生效。 第三步，计算损失。将预测结果和真实标签传入损失函数，得到损失标量。 第四步，反向传播。调用损失的backward()方法，自动计算所有可训练参数的梯度。 第五步，参数更新。调用优化器的step()方法，根据梯度和学习率更新模型参数。 整个循环不断重复，遍历完所有训练批次，就完成了一轮训练。
22.3.3 验证流程的注意事项
验证阶段和训练阶段有三个关键区别，必须严格遵守，否则验证结果会失真： 第一，切换模型模式。调用model.eval()进入评估模式，Dropout 会关闭，BatchNorm 使用全局统计量，保证输出稳定。 第二，关闭梯度计算。用torch.no_grad()上下文管理器包裹验证代码，不构建计算图，既节省显存，又提升速度。验证不需要反向传播，完全不需要梯度信息。 第三，不需要更新参数。验证只做前向传播和指标计算，不调用反向传播和优化器更新。 验证结束后，要调用model.train()切回训练模式，再开始下一轮训练。很多初学者验证后忘记切回训练模式，会导致后续训练效果异常差，这是非常高频的易错点。
22.3.4 指标记录与早停机制
训练过程中需要记录每一轮的训练损失、验证损失、各项评估指标，便于后续分析训练曲线。通常用列表保存每轮的指标值，训练结束后可以绘制损失曲线，直观观察收敛情况。 早停是工业界标配的正则化手段：连续多轮验证指标没有提升时，提前终止训练，防止过拟合。实现逻辑很简单：记录最佳验证指标和连续不提升的轮数，每轮验证后更新最佳值，超过耐心值还没提升就终止训练。早停不需要修改模型结构，几乎没有额外成本，是性价比最高的正则化方法之一。
22.4 GPU 加速：硬件算力的高效利用
深度学习训练的计算量极大，仅靠 CPU 往往需要数天甚至数周，GPU 的并行计算能力可以将训练速度提升几十上百倍。PyTorch 对 GPU 做了深度优化，只需要简单的设备迁移，就能让所有运算运行在 GPU 上。
24.2 设备迁移的核心逻辑
GPU 运算的核心原则是：参与运算的所有张量和模型，必须位于同一个设备上。CPU 张量不能直接和 GPU 张量运算，否则会报错。 设备迁移有两种常用写法：
.to(device)方法：张量和模型都支持这个方法，指定目标设备即可。模型调用.to(device)会把所有参数都迁移到对应设备，张量调用则返回新的 GPU 张量。.cuda()方法：直接迁移到 GPU，是早期的常用写法，兼容性好。
推荐统一使用device变量控制设备，代码开头定义设备：
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")这样同一份代码可以自动适配有 GPU 和无 GPU 的环境，可移植性更强。 需要特别注意：模型是 in-place 迁移，调用.to(device)后模型本身就被移动了；张量则是返回新对象，需要重新赋值接收。这是初学者很容易踩的坑，经常出现张量没赋值导致依然在 CPU 上运行的问题。
22.4.2 显存占用与基础优化
GPU 的显存是稀缺资源，大模型训练中显存往往是最大的瓶颈。基础的显存优化技巧包括： 第一，合理调整批次大小。批次大小是影响显存的最主要因素，显存不足时优先调小批次。 第二，及时释放无用张量。不再使用的中间变量可以手动删除，配合torch.cuda.empty_cache()释放缓存，不过这只是辅助手段，不能从根本上解决显存不足。 第三，验证阶段关闭梯度。关闭梯度计算能节省大量显存，尤其是大模型推理时，差异非常明显。 第四，使用更低的数值精度。比如从 float32 换成 float16，显存占用直接减半，这也是混合精度训练的核心思路。
22.4.3 多 GPU 训练基础
单卡显存不足时，就需要多 GPU 分布式训练。PyTorch 提供了两种基础的多卡方案：nn.DataParallel：数据并行的简单实现，单进程多线程，在主卡上做梯度汇总。优点是代码改动极小，只需要把模型包裹一层即可；缺点是效率不高，主卡容易成为瓶颈，不适合大规模多卡训练。DistributedDataParallel：多进程分布式训练，每个卡一个独立进程，同步梯度，效率更高，是工业界的标准方案。大模型的分布式训练、LoRA 微调，基本都基于 DDP 实现。 入门阶段先掌握单卡训练流程，多卡分布式是在此基础上的扩展，核心训练逻辑完全一致。
22.5 模型的保存与加载
训练好的模型权重需要持久化到硬盘，用于部署、断点续训、迁移学习。PyTorch 有两种保存加载模式，分别对应不同场景。
22.5.1 保存与加载 state_dict
推荐的标准方式是只保存模型的参数字典，也就是state_dict。它只包含权重和偏置等可训练参数，不包含模型结构代码，文件体积小，灵活性高。 保存：
torch.save(model.state_dict(), "model_weights.pth")加载时需要先实例化模型结构，再把权重加载进去：
model = MyModel() model.load_state_dict(torch.load("model_weights.pth"))这种方式要求加载时模型结构和保存时完全一致，否则会报错。它的优势是结构和权重分离，方便修改模型后加载部分权重，迁移学习、微调场景都用这种方式。大模型的权重文件，本质都是 state_dict 格式。
22.5.2 完整模型的保存与加载
第二种方式是直接保存整个模型，包含结构和权重：
torch.save(model, "full_model.pth")加载时直接加载：
model = torch.load("full_model.pth")这种方式代码简单，但灵活性差，依赖模型类的定义路径，代码目录变动就可能加载失败，而且文件体积更大。只适合简单的快速验证场景，正式项目不推荐使用。
22.5.3 断点续训：保存训练状态
完整的训练断点不仅要保存模型权重，还要保存优化器状态、当前轮次、最佳指标等信息，保证加载后可以接着训练，不需要从头开始。 通常把所有信息打包成一个字典保存：
checkpoint = { "epoch": epoch, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "best_loss": best_loss, } torch.save(checkpoint, "checkpoint.pth")加载时分别恢复模型、优化器、训练进度。长周期的大模型训练必须支持断点续训，避免训练中断导致前功尽弃，这是工程化训练的标配功能。
22.5.4 部分权重加载与严格匹配
迁移学习、微调场景下，经常需要加载预训练权重，但模型结构和预训练模型不完全一致。这时候可以设置strict=False，只加载名称匹配的权重，不匹配的自动忽略。
model.load_state_dict(pretrained_weights, strict=False)这是微调非常常用的技巧。大模型微调时，新增的下游任务头没有预训练权重，就用非严格加载，只加载主干网络的预训练参数。
22.6 工业级训练优化技巧
22.6.1 学习率调度器
固定学习率不是最优选择，训练过程中动态降低学习率，能够在前期快速收敛，后期精细调整，最终得到更优的结果。PyTorch 的torch.optim.lr_scheduler模块提供了多种调度策略。 最常用的有三类： StepLR：固定间隔按比例衰减学习率，简单直接，适合简单任务。 ReduceLROnPlateau：监控验证指标，指标不再提升时自动降低学习率，非常智能，不需要手动指定衰减步长，是工业界常用方案。 CosineAnnealingLR：余弦退火，学习率按照余弦曲线下降，配合预热使用效果极佳，是大模型训练的标准调度策略。 调度器的调用时机通常是每轮训练结束后调用step()，ReduceLROnPlateau 需要传入当前验证指标。
22.6.2 梯度裁剪
训练深层网络尤其是循环神经网络时，容易出现梯度爆炸，导致损失震荡甚至发散。梯度裁剪是最简单有效的解决方案：设置一个最大梯度范数，当梯度的范数超过阈值时，等比例缩小梯度，把梯度限制在安全范围内。
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)这行代码加在反向传播之后、参数更新之前即可。梯度裁剪几乎不增加计算量，却能极大提升训练稳定性，是训练深层网络的标配技巧。大模型训练中，梯度裁剪也是标准配置。
22.6.3 混合精度训练
默认的 float32 精度虽然稳定，但显存占用大、计算速度慢。混合精度训练将部分运算用 float16 半精度执行，部分保留 float32，在几乎不损失精度的前提下，显存占用减半，训练速度提升近一倍。 PyTorch 提供了原生的自动混合精度模块torch.cuda.amp，通过 GradScaler 缩放损失值，避免半精度下梯度过小下溢为零。 混合精度是大模型训练的必备优化，几乎所有大模型框架都默认开启。它不需要修改模型结构，只需要少量代码改动，性价比极高。
22.6.4 固定随机种子
深度学习训练有很多随机因素：权重初始化、数据打乱、Dropout、数据增强等，都会导致每次训练结果有差异。为了实验可复现，必须固定所有随机种子，包括 Python 随机数、NumPy 随机数、PyTorch CPU 和 GPU 随机数。 固定种子是科学实验的基本要求，保证相同的代码和数据，每次运行都能得到完全一致的结果。正式实验必须设置固定种子，否则结果不具备可复现性。
22.7 实战：完整的图像分类训练工程
结合本章所有知识点，可以搭建一个完整的工程化训练脚本，包含数据加载、训练验证循环、GPU 加速、学习率调度、早停、模型保存全流程。 整个脚本分为几个模块：配置参数定义、数据集与数据加载器构建、模型实例化与设备迁移、损失函数与优化器、学习率调度器定义、训练函数、验证函数、主训练循环。 主循环中逐轮训练和验证，记录损失和准确率，更新学习率，判断是否保存最佳模型，触发早停则终止训练。最后保存最终模型和训练日志。 这一套训练范式是通用的，无论是简单的图像分类，还是复杂的大模型微调，核心结构都是完全一致的，只是模型和数据集的具体实现不同。掌握了这个标准流程，就能快速迁移到任意深度学习任务中。
22.8 本章小结
本章系统讲解了 PyTorch 工程化训练的全流程，从数据加载到训练闭环，从硬件加速到权重管理，覆盖了工业级训练的核心知识点。核心内容回顾：
- Dataset 定义数据访问接口，DataLoader 负责批量调度加载，collate_fn 处理变长数据，构成了标准化的数据供给体系；
- 训练循环遵循梯度清零、前向传播、计算损失、反向传播、参数更新的五步标准流程；
- 验证阶段必须切换评估模式、关闭梯度，保证结果准确可靠，早停是性价比极高的正则化手段；
- GPU 加速需要统一设备，显存优化是大模型训练的核心课题，混合精度是高效的优化方案；
- 推荐保存 state_dict 格式的权重，支持断点续训和部分加载，适配微调与迁移学习场景；
- 学习率调度、梯度裁剪、固定随机种子是工业级训练的标配技巧，提升训练稳定性与最终效果。