PyTorch 使用 torch.nn.DataParallel 来实现基于数据并行的单机、多 GPU 并行训练。使用 DP 方式在多 GPU 上训练模型,需要保证模型能够在每个 GPU 上放得下,训练过程中会把整个模型都复制到每个 GPU 内,通过数据并行的方式提高训练模型的速度。 虽然 DataParallel 使用起来非常容易,但是通常不能够获得最佳的性能。DataParallel 在每一轮的前向传播过程中,会复制一遍模型,同时这种基于单进程、多线程的并行方式也会存在 GIL(Global Interpreter Lock) 竞争问题。 DP 数据并行训练流程 下面我们分析一下 DP 数据并行模式在多 GPU 的情况下训练模型的基本流程,如下图所示: 基于 DP 模式,模型训练的基本过程分为三个阶段,描述如下: 前向传播计算过程 1.Scatter mini-batch inputs to GPUs 我们通过指定 batch_size 大小,对输入的训练数据集进行了分割,在训练神经网络模型的前向传播过程中,会将每一个小批数据(Mini-Batch)分发到每一个 GPU 上。具体过程是:将 4 个小批数据 i1 ~ i4 复制到 GPU-1 上,再将 4 个小批数据分别发送到 GPU-1 ~ GPU-4