怎么看batchsize神经网络（增大Batch训练神经网络）

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

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

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

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

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

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

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

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

怎么看batchsize神经网络（增大Batch训练神经网络）(1)

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

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

怎么看batchsize神经网络（增大Batch训练神经网络）(2)

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

怎么看batchsize神经网络（增大Batch训练神经网络）(3)

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使用不均衡的问题，下图给出了相应解释：

怎么看batchsize神经网络（增大Batch训练神经网络）(4)

Forward and Backward passes with torch.nn.DataParallel

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

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

怎么看batchsize神经网络（增大Batch训练神经网络）(5)

Number of elements in the output of a language model

现有如下假设：数据集共含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的内部情况：

怎么看batchsize神经网络（增大Batch训练神经网络）(6)

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

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

分布式训练

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

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

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

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

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

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

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

怎么看batchsize神经网络（增大Batch训练神经网络）(7)

The main server (server 1) has an accessible IP and an open port for communication.

升级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并行。

怎么看batchsize神经网络（增大Batch训练神经网络）(8)