Lightning框架下CPU存数据集GPU训练的显存溢出解决方法

针对你遇到的数据集提前加载到GPU导致显存耗尽的问题，调整以下Trainer参数及配套配置即可解决：

• 关闭自动全量数据转移：不要设置Trainer(auto_move_data=True)，这个参数会强制把整个数据集一次性移到GPU，Lightning默认行为就是逐批次将数据从CPU传输到GPU，去掉这个配置就能恢复默认逻辑。
• 选择合适的分布式策略：如果用多GPU训练，不要使用strategy="dp"（数据并行），改用strategy="ddp"或strategy="ddp_spawn"。DP会把完整数据集拷贝到每个GPU，而DDP是每个进程独立加载自己的数据子集，大幅降低单GPU的数据显存占用。单GPU场景下确保devices=1即可。
• 调整批次相关参数：直接减小batch_size，如果需要保持等效训练批次，可以搭配Trainer(accumulate_grad_batches=N)，让模型累积N个批次的梯度再更新，既降低单批次显存占用，又不影响训练效果。
• 启用梯度检查点（可选）：如果模型本身显存占用高，添加Trainer(gradient_checkpointing=True)，通过牺牲少量计算速度来减少中间激活值的显存消耗，间接缓解数据加载带来的显存压力。

配套DataLoader优化（非Trainer参数但关键） #

确保你的DataLoader和Dataset遵循CPU预处理逻辑：

• DataLoader设置pin_memory=True，加速CPU到GPU的数据传输，但不要在Dataset的__getitem__方法中将数据移到GPU，所有数据预处理都在CPU完成，Lightning会自动在训练时将批次数据传输到GPU。

内容的提问来源于stack exchange，提问作者Bentium