PyTorch 使用 DP 模式实现数据并行

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