第一次用 Python 训练一个真正“大数据集”的时候我以为把几十万条样本一次性加载进内存是天经地义的事。结果程序跑了十几分钟内存占用一路飙升最后直接在终端里吐出一串 MemoryError整个进程当场没了。后来我才意识到这根本不是内存不够的问题而是我还没弄明白“分批训练 create batch”这件事——把数据切成一块一块让模型按批次去学这才是大数据集训练的正道。这篇文章我想把“分批训练”这件事从原理到实操彻底讲清楚。无论你是刚接触深度学习、正在跑 Kaggle 比赛还是做工业级的数据建模最终都会面对“数据量大到放不进内存/显存”的时刻。学会建造高效的 batch 流水线、选择合适的 batch_size、避开分批训练里那些隐蔽的坑比换个更大的服务器划算得多也更能体现一个工程师的基本功。1. 为什么大数据集不能“一口吃下”显存、内存与梯度更新的三重约束1.1 一个算账的例子50万张图片到底占多少内存我们先算一笔最简单的账。假设手头有一个图像数据集50万张 RGB 图片每张缩放到 224×224。单张图片的像素数量是 224×224×3约 15 万个数。如果以 float32 存储每个数占 4 字节一张图就是约 600 KB。50 万张就是 300 GB 左右。哪怕你把数据压成 uint8每像素 1 字节也要约 75 GB。普通开发机的内存通常在 16GB 到 64GB服务器稍高一些但也很难轻松吞下 300GB 的原始数据。更何况深度学习训练过程中真正吃内存/显存的大头往往不是原始数据而是模型前向传播时保存的中间激活值。输入一张 224×224 的图过一层卷积特征图可能比原图还大中间变量的体积轻松超过输入本身。所以“把整个数据集装进内存再开始训练”这条路线在规模上来之后基本是走不通的。1.2 全量梯度 vs 小批量梯度从数学上理解为什么要分批从优化算法的角度看典型的梯度下降有三种形态。一是批量梯度下降Batch Gradient Descent每次用全部样本计算梯度然后更新一次参数。这种方法数学上最稳定、梯度方向最准确但有两个致命问题内存/显存扛不住每次更新都极其缓慢几十万样本算一次梯度可能要几分钟甚至更久。二是随机梯度下降SGD每次只拿一条样本算梯度就更新。它内存占用极小但每次梯度方向噪声极大训练过程振荡严重收敛效率很差而且无法利用现代 GPU 的并行计算能力——一次只喂一条数据GPU 的计算单元大部分时间在空转。三是小批量梯度下降Mini-batch Gradient Descent也就是我们说的分批训练。每次取一“批”比如 32 条、64 条、128 条样本计算梯度然后更新参数。它既能把内存/显存占用控制在可接受范围又能利用 GPU 并行能力梯度方向也在“噪声”和“准确性”之间取了平衡。PyTorch 里的 DataLoader 本质上就是为这个“小批量梯度下降”服务的。1.3 分批之后模型到底在做什么一个 epoch 内部的执行顺序一个 epoch把所有训练数据完整过一遍内部训练循环的执行顺序是这样的DataLoader 按 batch_size 把数据集切成若干个小批次每个批次是若干条样本的集合。训练循环每次取出一个小批次把它送到 GPU 上。模型做一次前向传播算出这个批次的预测值和损失。通过反向传播算出损失对每个参数的梯度。优化器拿着梯度去更新一次模型参数。回到第 2 步取出下一个小批次继续这个过程直到所有批次都被处理完。所以“一个 epoch”其实是“把 2、3、4、5 步重复了 N 次”N 约等于“总样本数 ÷ batch_size”。比如 5 万条样本、batch_size 为 100一个 epoch 就是 500 次参数更新。这样每次更新的数据量不大但足够模型逐步学到规律。2. 分批执行的标准载体Dataset 与 DataLoader 的正确打开方式2.1 自定义 Dataset三个方法一个都不能少在 PyTorch 里DataLoader 负责“分批”但它取数据的逻辑依赖于一个 Dataset 对象。很多人一上来就抄网上的代码用 torchvision.datasets.ImageFolder 或者 TensorDataset一旦遇到自定义格式的数据就懵了。其实自定义 Dataset 只需要实现三个方法__init__、__len__、__getitem__。from torch.utils.data import Dataset import numpy as np import pandas as pd class MyTabularDataset(Dataset): def __init__(self, csv_path): # 这里不加载全部特征只加载文件路径或行索引 self.df pd.read_csv(csv_path, usecols[file_path, label]) self.file_paths self.df[file_path].tolist() self.labels self.df[label].tolist() def __len__(self): # 返回数据集总长度DataLoader 靠它知道一共有多少条 return len(self.file_paths) def __getitem__(self, idx): # 给定索引返回一条样本(特征, 标签) # 真正费时间的读取操作放在这里而不是 __init__ 里 data np.load(self.file_paths[idx]) label self.labels[idx] return data.astype(np.float32), label这里有一个新手容易犯的错误在__init__里把所有数据一次性读进内存。这等于绕开了分批的意义。正确做法是在__getitem__里按索引读取DataLoader 拿到索引后才会真正触发文件读取。2.2 DataLoader 参数batch_size、shuffle、drop_last 的底层含义Dataset 只是定义了“数据集长什么样”真正干分批活的是 DataLoader。它的几个核心参数每一个都不白给参数作用说明batch_size每次取多少条样本打包成一个 batch核心参数直接影响显存占用和训练速度shuffle每个 epoch 开始前是否打乱样本顺序设为 True 用于训练集防止样本顺序带来的偏见drop_last最后剩余不足一个 batch 的样本是否丢弃设为 True 可以避免最后一个 batch 过小num_workers用几个子进程来加载数据影响数据读取速度详见后文pin_memory是否把数据放到锁页内存加速向 GPU 拷贝用 GPU 训练时建议设为 Truecollate_fn自定义“怎么把样本打包成一个 batch”处理变长序列时必备一个容易被忽略的点是shuffle的真面目。DataLoader 内部其实是靠 Sampler 来决定取数顺序的。shuffleTrue等价于使用RandomSamplershuffleFalse等价于使用SequentialSampler。也就是说控制“按什么顺序取样本”的钩子在 Sampler 层DataLoader 只是把 Sampler 输出的索引拿去问 Dataset 要数据。后面讲进阶时你会看到采样器的威力。2.3 从“读一批数据”到“喂给模型”的完整链路数据从磁盘到 GPU 上的完整链路是Sampler 输出一批索引比如 [3, 17, 9, ...]。DataLoader 拿着索引调用Dataset.__getitem__得到一条条样本。collate_fn把这些样本堆叠成张量默认是把 numpy 数组转换成 torch.Tensor 并加一个 batch 维度。如果设置了pin_memoryTrue数据会先进入锁页内存。训练代码调用batch.to(device)数据从内存拷贝到显存。模型接收这个 batch开始前向传播。理解了这条链路排查问题就有方向了。比如训练时 GPU 利用率忽高忽低说明瓶颈可能在第 1 到第 5 步之间如果 GPU 利用率一直很高但显存爆了说明 batch_size 或者网络结构有问题。分段定位不要一上来就怀疑代码。3. batch_size 不是拍脑袋定的从显存、收敛速度与训练稳定性三个维度找答案3.1 显存估算公式与实测方法很多人问我“batch_size 设多少合适”我给的标准答案都是先看你显存能装下多少再结合收敛情况调。显存占用有几个固定来源模型参数、优化器状态Adam 还需要额外存一阶和二阶动量、前向传播的中间激活值、反向传播的梯度。其中激活值和梯度都跟 batch_size 成正比。粗略估算的方法是先用脚本随机生成一个和真实形状一致的 dummy batch然后加显存监控看差分峰值。更快的实测方法是写一个只跑一个 forward 和 backward 的脚本import torch import torch.nn as nn def estimate_batch_size(model, sample_feature, sample_label, max_memory_mb4096): model model.cuda() torch.cuda.reset_peak_memory_stats() try: # 逐步增大 batch for batch_size in [8, 16, 32, 64, 128, 256]: features sample_feature.repeat(batch_size, 1).cuda() labels sample_label.repeat(batch_size).cuda() optimizer torch.optim.Adam(model.parameters()) optimizer.zero_grad() out model(features) loss nn.CrossEntropyLoss()(out, labels) loss.backward() optimizer.step() peak_mb torch.cuda.max_memory_allocated() / 1024 / 1024 print(fbatch_size{batch_size}, peak_memory{peak_mb:.1f}MB) except RuntimeError as e: print(fbatch_size{batch_size} 时 OOM报错{e}) # 临时用 dummy 数据测即可无需真实数据集 class DummyNet(nn.Module): def __init__(self): super().__init__() self.fc nn.Sequential( nn.Linear(1024, 2048), nn.ReLU(), nn.Linear(2048, 1024), nn.ReLU(), nn.Linear(1024, 10) ) def forward(self, x): return self.fc(x)实际操作时我会把max_memory_mb设为显卡总显存的 70% 左右留一些余量给 CUDA 上下文和其他开销。这样测出来的 batch_size 就是“当前模型结构下比较稳的上限”。3.2 batch 越大收敛越好吗过大 batch 的隐性代价显存允许的情况下是不是 batch 越大就一定越好不是。大 batch 有几个隐性代价第一泛化能力会变差。实验和理论都表明过大的 batch_size 会让模型倾向于收敛到“尖锐极小值”测试集表现相对较差而适中偏小的 batch 由于梯度噪声的存在更容易逃出尖锐区域泛化更好。第二收敛曲线会更“僵硬”。大 batch 的梯度方向很平均loss 曲线一开始很好看但后期容易卡在某个平台期不动因为它失去了小 batch 那种“扰动带来的探索能力”。第三优化器状态的内存开销会被放大。尤其当你用 Adam 这类带动量状态的优化器时每个参数要额外存两份状态模型越大、batch 越大显存压力成倍增加。所以我的经验是batch_size 最好在 32 到 128 之间起步配合学习率一起调。如果你的显卡很豪横也不要一开始就上 512先跑 128 看收敛趋势再决定要不要放大。3.3 与学习率的配合一个实用调参经验这里有一个很多教程没讲透的点batch_size 变了学习率大概率也要跟着变。原因是梯度不再是一个“全量平均”而是一个 mini-batch 的估计。batch 越大梯度估计越稳定就可以用更大的学习率batch 越小梯度噪声越大学习率就得保守一些。一个常用的经验法则是线性缩放规则batch_size 从 B1 放大到 B2学习率大致按 B2/B1 倍进行缩放。比如原来用 batch32、lr0.001改成 batch64 后可以把 lr 调到 0.002 左右然后观察头几个 epoch 的 loss 是否稳步下降。如果不稳定就退回 0.0015。但线性缩放有个前提必须配合 warmup。因为大学习率在一开始很容易把参数推到远离收敛的位置前几个 epoch 先让学习率从一个很小的值线性升到目标值能有效避免震荡。我自己的习惯是前 5 个 epoch 做 warmup之后用余弦退火把学习率慢慢降下来。这套组合在大 batch 场景下非常稳。4. 分批训练必踩的三个坑尾部 batch、样本顺序和随机种子4.1 最后一个 batch 太小要不要 drop_last假设训练集有 1003 条样本batch_size 是 100那么前 10 个 batch 都是满的 100 条最后一个 batch 只有 3 条。这个过小的尾部 batch 会带来两个问题一是如果模型里有 BatchNorm 层小 batch 上计算出的均值和方差非常不稳定几十条还算好几条就会给 BN 层注入大量噪声影响训练稳定性。二是尾部 batch 的梯度只来自 3 条样本在整个训练过程中占比很微弱但它会让“一个 epoch 到底更新几次参数”变得不规律不好复盘。我的建议是训练集一律drop_lastTrue把尾部那一点点样本丢掉。不用心疼——数据足够大时几十条样本的影响几乎可忽略。但如果你的数据很小比如几千条丢掉太多会导致每个 epoch 利用的数据变少这时可以考虑把batch_size调整成能被样本总数整除的数或者实现一个不均匀切分策略。验证码、时间序列这类要求“每条样本都必须用上”的任务可以把drop_lastFalse但单独处理尾部 batch例如给尾部 batch 一个更小的学习率更新权重避免它扰乱整体参数更新。4.2 shuffle 为什么如此重要一个反例如果你的训练集是按标签顺序排好的比如前 10 万条都是类别 A中间 10 万条都是类别 B后面都是类别 C而你忘了打开shuffleTrue那模型学到的会是灾难性的结果。前几个 epoch 里模型看到所有 batch 全是同一类样本梯度方向被某个单一类别主导参数被推得偏得离谱。就算数据不是刻意按标签排序很多数据集也存在潜在的顺序偏差。比如按时间顺序采集的数据早批次的样本和晚批次的样本分布可能不一样。shuffle 的作用就是让每个 batch 尽量代表整体数据分布避免模型学习到“顺序”这个伪特征。所以数据加载的正确姿势是训练集用shuffleTrue验证集/测试集用shuffleFalse。验证集不需要打乱因为验证集只做前向传播不打乱反而方便和标签原始顺序对齐做分析。4.3 随机种子和可复现性魔改 DataLoader 之后的结果对不上分批训练引入了一个额外随机性来源数据加载的顺序和getitem里可能的随机操作。比如你的__getitem__里有数据增强每次调用都产生随机扰动又比如shuffleTrue时每个 epoch 的取数顺序不同。这些随机性本身没问题但你做实验对比不同超参时如果每次跑的随机性都不一样很难判断效果差异到底是来自参数改动还是偶然波动。所以项目一开始就要固定随机种子。推荐在训练脚本最前面统一设置import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False set_seed(42)这里有两个容易被忽略的细节torch.backends.cudnn.deterministicTrue会把 cuDNN 的卷积算法固定下来让卷积结果可复现但会损失一些性能torch.backends.cudnn.benchmarkFalse则避免算法选择带来的不确定性。如果你只是做日常调试可以不开 deterministic如果是写论文级的对比实验这两个开关必须打开。我自己的经验是加了这两个开关之后同一套代码跑两次loss 曲线能高精度重合排查问题会省很多事。5. 数据加载速度才是训练效率的隐藏瓶颈num_workers 与 pin_memory 实测记录5.1 num_workers不是越大越快很多人认为num_workers设成 CPU 核心数越高越好其实不一定。num_workers表示用几个独立子进程去预取数据。数据加载是一个典型的 IO 密集型任务瓶颈通常在磁盘读取速度和解码计算上而不是进程数。我实测过一组数据在一台 16 核的机器上同样训练一个图像分类模型num_workers0主进程同步加载时 GPU 利用率只有 40% 左右调到 4利用率跳到 85%调到 8利用率到了 92%再往上调到 16利用率反而没涨多少内存却额外多了好几个 GB。原因在于每个 worker 子进程都会启动一个 Python 解释器并且会复制部分内存状态。worker 数量超过某个阈值后CPU 切换到进程的开销和内存复制消耗开始超过收益。即使数据集较大num_workers设置在 CPU 核心数的一半到四分之三之间往往就是性价比最高的区间。代码里通常可以写成min(os.cpu_count(), 8)先保证能跑再根据 GPU 空闲率微调。5.2 pin_memory 到底有没有用pin_memory这个参数有点“让人又爱又恨”。它的作用是把数据放进锁页内存使得 CPU 到 GPU 的拷贝可以走高速 DMA 通道。如果用 GPU 训练建议直接设为pin_memoryTrue一般能带来 10% 到 30% 的数据传输提速代价是锁页内存比普通内存贵占用过多时会影响系统整体性能。如果不用 GPU或者训练数据量小、显存完全无压力那么这个参数基本没意义。还有一种情况需要注意如果num_workers0pin_memoryTrue的效果会打折扣因为数据传输和训练计算被串行化高速通道的优势被掩盖了。最推荐组合是num_workers 0且pin_memoryTrue这样数据预取和传输可以并行。5.3 实践中如何判断“数据加载卡脖子了”分管数据加载和模型计算的部门如果配合不好GPU 经常会等数据。怎么判断最直观的办法是看 GPU 利用率训练时在终端输入nvidia-smi观察 Utilization或者用watch -n 1 nvidia-smi动态观察如果利用率长期低于 70%而显存占用并没有打满多半就是数据加载拖了后腿。你也可以在训练循环里手动计时。比如对比下面两段代码的执行时间# 情况A每步消耗的时间 import time start time.time() for batch in train_loader: # 模型前向反向更新 ... end time.time() print(fstart to first batch: {end - start:.2f}s) break如果发现从循环开始到拿到第一个 batch 耗时很长或者中途某几步耗时明显大于其他步都是数据加载不平衡的信号。还有一种常见现象是每个 epoch 刚开始时 GPU 利用率骤降过两三秒才回升——这是因为上一个 epoch 结束时 DataLoader 的预取缓存被清空了新一批数据还没进来。这种情况下优先检查num_workers是否太低然后是pin_memory接着看磁盘是否太慢机械硬盘换成 SSD 能解决很多问题。实在不行还可以用数据预处理管线把图片离线缩放到小尺寸减少每个 batch 的加载压力。6. 再往深处走采样器、梯度累积与变长 batch6.1 Sampler控制“取哪一批”的钩子前面说过DataLoader 的shuffle参数本质是切换 Sampler。如果你有更复杂的取批需求需要直接操作 Sampler。最常用的场景是类别不平衡。比如某些任务中正样本只有 1%负样本 99%直接用随机采样会让每个 batch 几乎全是负样本模型学不到正样本的特征。这时可以用WeightedRandomSampler给样本分配权重让每个类别在采样时被抽中的概率尽量均衡from torch.utils.data import WeightedRandomSampler # labels 是全体训练集标签 class_counts np.bincount(labels) class_weights 1.0 / class_counts sample_weights [class_weights[label] for label in labels] sampler WeightedRandomSampler(sample_weights, num_sampleslen(labels), replacementTrue) dataloader DataLoader(dataset, batch_size64, samplersampler)注意WeightedRandomSampler一般需要设置replacementTrue也就是允许同一轮 epoch 中重复抽取某些样本否则某些小类样本的期望出现次数很难保证。用了这种采样器之后一个 epoch 的含义变成了“总共抽取 num_samples 次”而不是“把数据集完整过一遍”日志里要留意。6.2 显存不够的曲线救国梯度累积假设你的显卡只够跑 batch_size16但根据经验你想实现 batch_size64 的效果。除了换卡还有一个办法梯度累积Gradient Accumulation。原理很简单连续 4 个 batch每个 batch 都算梯度但不立即更新参数每攒够 4 次把梯度累加起来做一次参数更新。这样每次更新的统计效果约等于 batch_size64但实际每个 step 只占用 16 条样本的显存。accumulation_steps 4 optimizer.zero_grad() for i, batch in enumerate(train_loader): features, labels batch outputs model(features) loss loss_fn(outputs, labels) loss loss / accumulation_steps # 缩放损失防止梯度累加后过大 loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这里有一个细节loss要除以accumulation_steps相当于做了一次平均否则累积梯度的量级会变成原来的 N 倍学习率得重新调。另外梯度累积对 BatchNorm 层有隐含影响——BN 的统计量依然是从 16 条样本的小 batch 里算出来的不会自动变成 64 条的水平。所以这种方法在纯全连接、无 BN 的模型上效果更贴近真实大 batch。6.3 自然语言数据的分批变长序列与 collate_fn处理文本数据时每条样本长度不一样没法直接堆叠成 batch。常规做法是 padding把同一 batch 内的文本补齐到相同长度。但一个常见失误是把整个数据集的所有文本都 pad 到全局最长长度——这样很浪费算力尤其当长文本只有几条时。正确思路是先按长度大致排序让相近长度的样本进入同一个 batch再在collate_fn里做动态 padding只 pad 到“当前 batch 的最大长度”。from torch.nn.utils.rnn import pad_sequence def collate_batch(batch): texts, labels zip(*batch) # texts 是若干 [seq_len] 的 tensor按当前 batch 最大长度补零 texts_padded pad_sequence(texts, batch_firstTrue, padding_value0) labels torch.tensor(labels) return texts_padded, labelscollate_fn的触发时机是在“取完一个 batch 对应索引的样本之后、把这些样本堆叠成张量之前”。所以它天然适合做变长处理、动态 pad、甚至可以把类别标签转成 one-hot。很多 NLP 的踩坑都出在默认collate_fn无法处理不同长度的样本上自定义之后问题迎刃而解。7. 一点个人经验与收尾分批训练这件事听起来只是“把数据切一切”真正做好之后训练全流程的稳定性和效率都会有明显提升。我自己踩过不少坑从最开始的内存爆炸到后来 GPU 利用率不满再到 BN 层因为尾部 batch 疯狂波动每一步都是在和“batch”这个概念较劲。最后分享一个排查习惯训练脚本开头我会先打印数据集总样本数、batch_size、每个 epoch 的 batch 数、num_workers等参数同时在终端开着watch -n 1 nvidia-smi观察显存和 GPU 利用率。任何一项不对劲都能在几分钟内定位到大方向。多花这几分钟后续训练可以省下几小时——这个买卖很划算。