deep-learning-for-image-processing 多 GPU 分布式训练实战指南：launch 与 spawn 两种启动方式详解

• 示例工程

本指南围绕 pytorch_classification/train_multi_GPU 目录，系统讲解在 deep-learning-for-image-processing 项目中基于 PyTorchtorch.distributed实现的多 GPU 图像分类训练方案。你将掌握torch.distributed.launch与torch.multiprocessing两种分布式训练启动方式的完整命令与参数含义，理解DistributedSampler、SyncBatchNorm、学习率倍增、跨进程指标归约等底层实现细节，并学会根据训练时间与准确率曲线评估多卡加速效果。

项目背景：为什么要在这个分类仓库中引入多 GPU 训练

在 train_multi_GPU 目录中，仓库以 ResNet34 在花卉数据集（flower_photos，默认 5 分类）上的训练为示范任务，同时提供了三个可以直接运行的脚本：

• train_single_gpu.py：单 GPU 基线训练脚本，用于与多卡结果做对照；
• train_multi_gpu_using_launch.py：通过torch.distributed.launch工具启动的多卡训练脚本；
• train_multi_gpu_using_spawn.py：通过torch.multiprocessing手动创建子进程启动的多卡训练脚本。

三个脚本共用同一条数据处理链路：utils.read_split_data负责按 8:2 划分训练集与验证集并生成class_indices.json，my_dataset.py 中的MyDataSet负责读取图像，而训练与验证的循环逻辑统一收敛在 multi_train_utils/train_eval_utils.py，分布式进程组的初始化则集中在 multi_train_utils/distributed_utils.py。这种"训练逻辑与启动方式解耦"的组织方式，让同一套代码可以随意切换启动入口。

一、方式一：使用 torch.distributed.launch 启动

1.1 基本启动命令

这是文档给出的标准启动方式，在仓库根目录下执行：

python -m torch.distributed.launch --nproc_per_node=8 --use_env train_multi_gpu_using_launch.py

其中：

• --nproc_per_node表示每台机器上并行的进程数，也就是实际使用的 GPU 数量（上例为 8，即单机 8 卡全量训练）；
• --use_env表示通过环境变量把 rank 信息传给训练脚本。启用后，torch.distributed.launch会自动设置RANK、LOCAL_RANK、WORLD_SIZE等环境变量，脚本内通过 distributed_utils.py 的init_distributed_mode读取。

1.2 指定使用某几块 GPU

如果机器上有 8 块 GPU 但只想用其中的第 1 块和第 4 块（索引从 0 开始），文档给出的命令是：

CUDA_VISIBLE_DEVICES=0,3 python -m torch.distributed.launch --nproc_per_node=2 --use_env train_multi_gpu_using_launch.py

这里的关键是CUDA_VISIBLE_DEVICES环境变量先对物理 GPU 做了一次"可见性过滤"，把物理卡 0 和 3 重映射为逻辑卡 0 和 1，因此--nproc_per_node只需写 2。该命令等价于"用第 1 块和第 4 块 GPU 各跑一个训练进程"，是日常多卡训练中最常用的 GPU 挑选写法。

1.3 launch 方式的底层初始化流程

torch.distributed.launch的工作机制可以概括为：它读取--nproc_per_node，在当前机器上为每一块卡派生一个训练子进程，并通过环境变量把分布式信息注入子进程。脚本收到这些环境变量后，在 train_multi_gpu_using_launch.py 调用init_distributed_mode，其内部逻辑为：

if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: args.rank = int(os.environ["RANK"]) args.world_size = int(os.environ['WORLD_SIZE']) args.gpu = int(os.environ['LOCAL_RANK'])

随后设置NCCL通信后端（NVIDIA GPU 官方推荐），调用dist.init_process_group建立进程组，并用dist.barrier()同步所有进程，确保各进程统一就绪后才开始训练。值得说明的是，init_distributed_mode还兼容了 SLURM 作业调度环境（通过SLURM_PROCID推导 rank），因此同一套脚本既能跑在裸机多卡环境，也能适配部分集群调度器。

1.4 命令行参数一览

脚本通过argparse暴露了完整的训练超参数，全部参数及默认值如下：

二、方式二：使用 torch.multiprocessing 手动启动

2.1 启动命令

不需要torch.distributed.launch包装器，直接运行即可：

python train_multi_gpu_using_spawn.py

该脚本在main之外自行完成进程编排：读取--world-size（默认 4，即单机使用 4 块 GPU），然后为每个 rank 创建独立的torch.multiprocessing.Process并逐个start()与join()，核心代码位于 train_multi_gpu_using_spawn.py：

world_size = opt.world_size processes = [] for rank in range(world_size): p = Process(target=main_fun, args=(rank, world_size, opt)) p.start() processes.append(p) for p in processes: p.join()

2.2 spawn 方式的底层初始化差异

与 launch 方式依赖外部注入环境变量不同，spawn 方式在main_fun内部手动完成进程组初始化：

os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "12355" args.rank = rank args.world_size = world_size args.gpu = rank args.distributed = True torch.cuda.set_device(args.gpu) args.dist_backend = 'nccl' dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) dist.barrier()

即：自行指定主节点地址MASTER_ADDR=localhost与端口MASTER_PORT=12355（单机场景），每个进程按传入的rank绑定到对应 GPU（args.gpu = rank），再调用dist.init_process_group建立进程组。这种写法把进程编排逻辑完全掌握在训练脚本手中，适合单机多卡或需要精确控制进程生命周期的场景；而 launch 方式则把进程编排交给 PyTorch 官方工具，更适合与集群工具链（如 SLURM）配合。

2.3 两种方式的等价内核

无论用哪种方式启动，进入训练主循环后逻辑完全一致：DistributedSampler切分数据、DDP 包装模型、SyncBatchNorm转换、学习率倍增、每 epoch 评估与保存权重。两种启动方式可以视为"进程编排层"的两种实现，训练核心只依赖init_distributed_mode/dist.init_process_group之后的统一代码路径，这也是仓库将 multi_train_utils 单独抽出的原因。

三、源码级原理：多卡训练中必须处理的五个细节

3.1 学习率按 GPU 数量倍增

多卡训练时全局批大小等于"每卡 batch_size × 卡数"，为保证收敛行为与单卡一致，脚本在训练开始前执行args.lr *= args.world_size（见 train_multi_gpu_using_launch.py）。例如默认--lr 0.001在 8 卡下会实际变为 0.008。这一点在切换卡数时务必留意，否则会改变训练动态。

3.2 DistributedSampler 保证数据不重叠且可复现

两个多卡脚本都使用torch.utils.data.distributed.DistributedSampler为每个进程切分样本，并用BatchSampler按batch_size打包（drop_last=True丢弃尾部不足一个 batch 的样本）。关键在于每个 epoch 开始前必须调用train_sampler.set_epoch(epoch)（见 train_multi_gpu_using_launch.py），否则每个 epoch 各进程拿到的样本划分顺序完全相同，数据 shuffle 失效。这是 PyTorch 分布式训练中最容易被遗漏的细节。

3.3 各进程权重初始化必须一致

当没有预训练权重时，脚本先由 rank 0 进程把初始化权重保存到系统临时目录（tempfile.gettempdir()下的initial_weights.pt），随后调用dist.barrier()等待所有进程就绪，再由各进程从临时文件加载同一份权重。源码注释特别提醒：加载时必须指定map_location=device，否则所有权重都会被加载到第一块 GPU 上，导致显存占用失衡（见 train_multi_gpu_using_launch.py）。

3.4 跨进程指标归约：reduce_value 与 all_reduce

训练损失和验证正确数都需要跨进程汇总，仓库在 distributed_utils.py 中实现了reduce_value：单卡直接返回原值；多卡时通过dist.all_reduce对所有进程的张量求和，average=True时再除以world_size得到均值。训练循环中的mean_loss就是先对每个 batch 的 loss 做reduce_value(loss, average=True)再滑动平均，验证阶段的sum_num则用average=False累加所有进程的正确预测数，最后除以val_sampler.total_size得到准确率。这也是 train_eval_utils.py 中train_one_epoch与evaluate的核心逻辑。

3.5 SyncBatchNorm：多卡训练 BN 的正确姿势

只有训练带 BN 结构的网络（如 ResNet34）时，启用 SyncBatchNorm 才有意义。脚本在args.syncBN为 True 时执行：

model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

它会遍历模型，把所有 BatchNorm2d 层替换为同步版本，使每个 BN 层的均值/方差统计在全局 batch（所有卡的样本）上计算，从而让"每卡 batch_size=16、8 卡并行"的效果接近"单卡 batch_size=128"。代价是训练更耗时（BN 统计需要跨进程通信），这一点在文档对应的 syncbn.png 曲线中可以得到验证。转换完成后再执行torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])包装为 DDP 模型，分布式梯度同步由 DDP 自动完成。

3.6 主进程职责：打印、TensorBoard 与权重保存

多卡环境下只有 rank 0 进程负责对外输出，避免重复打印与并发写文件。脚本中凡是if rank == 0:的代码块，都只做以下三件事（见 train_multi_gpu_using_launch.py）：

1. 打印当前 epoch 的验证准确率；
2. 向 TensorBoard 写入loss、accuracy、learning_rate三条标量曲线（启动命令为tensorboard --logdir=runs，浏览器访问http://localhost:6006/）；
3. 保存model.module.state_dict()（注意 DDP 模型需要取.module才能拿到原始模型的权重）到./weights/model-{epoch}.pth。

训练结束后 rank 0 还会清理临时权重文件，并调用cleanup()（即dist.destroy_process_group()）销毁进程组，完成分布式环境的正常退出。

四、训练效果验证：时间与精度对比

为了量化多卡收益，仓库在 plot_results.py 中提供了三组绘图数据，分别对应文档中的三张效果图。

4.1 训练时间对比

根据 plot_results.py 中记录的示例数据，单 epoch 训练耗时随 GPU 数量增加明显下降：1 卡约 9 秒、2 卡约 5.5 秒、4 卡约 3 秒、8 卡约 2 秒。可以看到 8 卡相比单卡在该示例任务上的加速效果显著，但由于多卡存在进程通信开销（尤其是启用 SyncBatchNorm 后），实际加速比通常不会达到理想的线性倍率，卡数越多边际收益越有限。

4.2 是否使用 SyncBatchNorm

图中对比了 30 个 epoch 内"多卡训练 + 未使用 SyncBatchNorm"与"多卡训练 + 使用 SyncBatchNorm"两条准确率曲线。从 plot_results.py 记录的示例数据看，使用 SyncBatchNorm 的曲线在训练后期准确率整体略高（例如第 29 个 epoch 约 0.81 对 0.80），说明在小 batch 多卡场景下，全局 BN 统计有助于稳定并提升最终精度。需要强调的是，这里展示的是项目作者在特定数据集与超参数下的单次实验示例，不同任务、数据集和卡数下结论可能不同，建议以自身实验为准。

4.3 单 GPU 与多 GPU 训练曲线

这张图把三条曲线放在一起：单 GPU 训练、多 GPU 不使用 SyncBatchNorm、多 GPU 使用 SyncBatchNorm。从 plot_results.py 记录的示例数据看，三者最终都能收敛到相近的准确率区间（约 0.80–0.83），验证了多 GPU 训练在保证精度的前提下显著缩短训练时间这一核心价值；同时也说明多卡只是"加速手段"，模型的最终精度仍取决于数据集、超参数与训练策略本身。

五、运行环境与前置条件

• 训练脚本会先检查torch.cuda.is_available()，无可用 GPU 时直接抛出EnvironmentError；
• requirements.txt 中列出的依赖版本为torch==1.13.1、torchvision==0.7.0、matplotlib==3.2.1、tqdm==4.42.1，实际运行时建议使用不低于此版本的 PyTorch，并确保多卡环境已正确安装 NVIDIA NCCL 通信库（分布式默认使用nccl后端）；
• 数据集需按"一个子目录一个类别"的组织方式放置，脚本在 utils.py 的read_split_data中自动按 8:2 划分，并校验训练集与验证集均非空；
• 若想对照单卡基线，可直接运行 train_single_gpu.py，其超参默认值与多卡脚本保持一致（--epochs 30、--batch-size 16、--lr 0.001、--lrf 0.1），便于直接对比训练曲线。

六、实践建议与注意事项

1. 卡数变化时务必检查学习率：多卡脚本会自动将--lr乘以world_size，单卡与多卡切换时学习率实际生效值不同，若追求完全一致的行为需要手动调整传入的--lr。
2. --syncBN默认开启：该参数默认值为 True。若数据量小、每卡 batch 足够大，或追求更短的训练时间，可以显式关闭（如--syncBN False），但要注意 argparse 对布尔参数的处理，推荐使用--syncBN 0或改为action='store_true'风格显式传入。
3. 验证集准确率的计算口径：多卡脚本用sum_num / val_sampler.total_size计算 acc，其中sum_num是经过all_reduce汇总的全局正确数；单卡脚本用sum_num / len(val_data_set)，两者口径一致（都是全局验证集），可直接对比。
4. 分布式训练的通用性：本目录的 multi_train_utils 工具模块（init_distributed_mode、cleanup、reduce_value、train_one_epoch、evaluate）与同仓库 mini_imagenet/multi_train_utils 等目录中的实现思路一致，这套"launch / spawn 双入口 + 统一训练工具"的代码组织方式可以直接迁移复用到其他分类网络（如 vision_transformer、swin_transformer）乃至目标检测、分割等更大规模模型的多卡训练中。

通过本指南，你已经掌握了 deep-learning-for-image-processing 仓库中多 GPU 训练的两种标准启动姿势、每个命令行参数的作用，以及分布式训练中最容易出错的采样器、权重同步、指标归约与 BN 统计等底层细节，可以直接在自己的多卡机器上复现上述训练流程并评估加速收益。

• 示例工程