增大Batch训练神经网络：单GPU、多GPU及分布式配置的实用技巧

摘要： 增大Batch训练神经网络：单GPU、多GPU及分布式配置的实用技巧

2018年中的大部分时间，我都在尝试利用训练神经网络克服GPUs的局限。无论是在包含1.5亿个参数的语言模型中，比如OpenAI’s huge Generative Pre-trained Transformer (or the recent and similar BERT model)，还是在拥有3000万个输入元素的神经网络中，我都只能利用GPU处理很少的训练样本。

可是若想利用随机梯度下降算法得出不错的结果，大批量的训练样本必不可少。

如果你的GPU只能处理少量样本，该如何训练大批量模型呢？

接下来，我将介绍几类工具和技巧。

本文主要会讨论PyTorch框架，并就以下几个问题进行探讨：

1. 当训练批量甚至单个训练样本大于GPU内存时，如何训练模型；
2. 如何高效地利用多GPU机器；
3. 如何在分布式设备上简单的使用多个机器。

在一个或多个GPU上训练大批量模型

你构建了一个不错的模型，可在尝试处理更多样本时，却得到CUDA RuntimeError：内存不足。

根据网友的回答你明白，加倍批量可以对结果进行优化。

此时，梯度累积（accumulating gradients）可以帮助到你。

PyTorch代码如下所示：

predictions = model(inputs)               # Forward pass
loss = loss_function(predictions, labels) # Compute loss function
loss.backward()                           # Backward pass
optimizer.step()                          # Optimizer step
predictions = model(inputs)               # Forward pass with new parameters

loss.backward()计算出每个参数的梯度，并存储在parameter.grad中。

梯度累积意味着，在调用potimizer.step()实现梯度下降之前，我们会求取parameter.grad张量中的几个反向操作的梯度和。

如下是使用梯度累积训练模型的示例。

model.zero_grad()                                   # Reset gradients tensors
for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)                     # Forward pass
    loss = loss_function(predictions, labels)       # Compute loss function
    loss = loss / accumulation_steps                # Normalize our loss (if averaged)
    loss.backward()                                 # Backward pass
    if (i+1) % accumulation_steps == 0:             # Wait for several backward steps
        optimizer.step()                            # Now we can do an optimizer step
        model.zero_grad()                           # Reset gradients tensors
        if (i+1) % evaluation_steps == 0:           # Evaluate the model when we...
            evaluate_model()                        # ...have no gradients accumulated

扩展

我们甚至可以在GPU上训练一个连样本都无法加载得模型，并且可以使用梯度检查点（gradient-checkpoingting）节省计算资源。
梯度检查点会将我们连续计算的元前馈和元反向传播切分成片段。但由于需要增加额外的计算以减少内存需求，该方法效率不高。不过，它在某些示例中又有较为明显的优势，比如在长序列上训练RNN模型，点击此处查看详情。

或有兴趣可进入下列文档进行查询：

TensorFlow：https://github.com/openai/gradient-checkpointing
PyTorch doc：https://pytorch.org/docs/stable/checkpoint.html

A “Memory-poor” strategy that needs O(1) memory (but requires O(n²) computation steps) — From Yaroslav Bulatov’s nice post: https://medium.com/tensorflow/fitting-larger-networks-into-memory-583e3c758ff9

多GPU机器

在多GPU服务器上训练PyTorch模型首选torch.nn.DataParallel。该策略能够在多个指定设备上按照batch dimension分割输入，实现并行化模块。

DataParallel实现如下所示：

parallel_model = torch.nn.DataParallel(model) # Encapsulate the model
predictions = parallel_model(inputs)          # Forward pass on multi-GPUs
loss = loss_function(predictions, labels)     # Compute loss function
loss.backward()                               # Backward pass
optimizer.step()                              # Optimizer step
predictions = parallel_model(inputs)          # Forward pass with new parameters

但DataParallel存在GPU使用不均衡的问题，下图给出了相应解释：

在前向传播的第四个步骤（见右上）中，GPU-1汇集了所有并行计算的结果。

通过下列所示的方式能够计算出语言模型输出的大小：

现有如下假设：数据集共含4万词汇，序列中包含250 tokens，每个batch 包含32个示例，每个元素4 bytes，模型的输出占用1.2GB。但我们需要2.4GB的内存才能存储相关的梯度张量。

这种存储方式会使得GPU-1被过度使用，从而造成GPU使用不均衡的问题。

多GPU机器上的负载均衡

想要解决GPU使用不均衡的问题需要将每部分输出都保留在原有的GPU上，而不汇集于GPU-1。

张航开源了名为PyTorch-Encoding的包，可用于缓解上述问题。

我对这个开源包做了一些调整，你可以点击此处下载parallel.py。此包中包含两个模块：DataParallelModel以及DataParallelCriterion，如下所示：

from parallel import DataParallelModel, DataParallelCriterion

parallel_model = DataParallelModel(model)             # Encapsulate the model
parallel_loss  = DataParallelCriterion(loss_function) # Encapsulate the loss function

predictions = parallel_model(inputs)      # Parallel forward pass
                                          # "predictions" is a tuple of n_gpu tensors
loss = parallel_loss(predictions, labels) # Compute loss function in parallel
loss.backward()                           # Backward pass
optimizer.step()                          # Optimizer step
predictions = parallel_model(inputs)      # Parallel forward pass with new parameters

DataParallelModel不同于torch.nn.DataParallel的是，前向传播的输出（predictions）没有汇集在GPU-1中，而是作为n_gup张量的元组分布在相应的GPU上。

DataParallelCriterion容器封装了损失函数，并且将n_gpu张量的元组和目标标签张量作为输入。

下图描述了DataParallelModel/DataParallelCriterion的内部情况：

下面有两个特殊情况，并给出了解决办法：

1. 模型输出了一些张量：你可以利用output_1,output_2 = zip(*predictions)分解它们。
2. 若你不想并行计算损失函数，则可以利用gathered_prdictions = parallel.gather(predictions)收集张量。

分布式训练

PyTorch中的DistributedDataParallel可以帮助我们在遇到大批量训练问题时，拥有控制多个服务器的运算能力。

但值得注意的是：由于对每个节点都要启动一个独立的Python训练脚本，在设定时需要注意改变工作流程。

每个脚本在训练中都会拥有：

1. 它自己的优化器，在每次迭代中都执行一个完整的优化，不需要参数传输。
2. 一个独立的Python解释器：能够避免GIL-freeze

在后面我们将通过代码进行讨论：

torch.distributed包能够为同步分布式运算提供低级原语，基于此构建得到DistributedDataParallel。你可以通过阅读文档以及教程对其进行进一步理解。

接下来，我们将使用具有两个4-GPU的服务器。

升级Python脚本以适用分布式训练

首先，我们需要对脚本进行升级，使其能够独立的在机器（节点）中运行。我们想要完全实现分布式，并且在每个结点的每个GPU上独立运行进程，这一共需要8个进程。

接下来，初始化分布式后端，封装模型以及准备数据，这些数据用于在独立的数据子集中训练进程。更新后的代码如下：

from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader

# Each process runs on 1 GPU device specified by the local_rank argument.
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')

# Encapsulate the model on the GPU assigned to the current process
device = torch.device('cuda', arg.local_rank)
model = model.to(device)
distrib_model = torch.nn.parallel.DistributedDataParallel(model,
                                                          device_ids=[args.local_rank],
                                                          output_device=args.local_rank)

# Restricts data loading to a subset of the dataset exclusive to the current process
sampler = DistributedSampler(dataset)

dataloader = DataLoader(dataset, sampler=sampler)
for inputs, labels in dataloader:
    predictions = distrib_model(inputs.to(device))         # Forward pass
    loss = loss_function(predictions, labels.to(device))   # Compute loss function
    loss.backward()                                        # Backward pass
    optimizer.step()                                       # Optimizer step

为Python脚本加载多个实例

现在，我们将在每个服务器上启动训练脚本的实例。

我们使用PyTorch中的torch.distributed.launch运行脚本。它能用于环境变量的设置，并使用正确的local_rank参数调用脚本。

最主要的是第一台机器，所有的机器都要求能对它进行访问。因此，它需要拥有一个可以访问的IP地址（示例中为：196.168.1.1）以及一个开放的端口（示例中为：1234）。我们将使用torch.distributed.launch在第一台机器上运行脚本，具体如下：

python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=1234 OUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of our training script)

同样在第二台机器中运行脚本：

python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr="192.168.1.1" --master_port=1234 OUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of our training script)

除了—node_rank参数之外，上述两个命令相同。

扩展

如果你觉得在计算机集群上运行一组几乎相同的命令有些枯燥，可点击此处了解GNU并行。

云服务器99元拼团购！拉新还可赢现金红包！300万等你瓜分！
马上一键开团赢红包： http://click.aliyun.com/m/1000019899/

---

本文作者：【方向】

阅读原文

本文为云栖社区原创内容，未经允许不得转载。