基于PyTorch使用大batch训练神经网络

最新推荐文章于 2024-03-01 00:51:39 发布

mingo_敏

最新推荐文章于 2024-03-01 00:51:39 发布

阅读量6.8k

点赞数 4

分类专栏： pytorch 文章标签： pytorch

本文链接：https://blog.csdn.net/shanglianlm/article/details/89437058

版权

pytorch 专栏收录该内容

85 篇文章 29 订阅

订阅专栏

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

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

当增大batch size时，经常会遇到CUDA RuntimeError：内存不足，但是为了提高性能，我们又需要增大batch size时，可以试试 梯度累积。
在这里插入图片描述
梯度下降优化算法的五个步骤。

与之对等的 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 中。

累积梯度意味着，在调用 optimizer.step() 实施一步梯度下降之前，我们会对 parameter.grad 张量中的几个反向运算的梯度求和。在 PyTorch 中这一点很容易实现，因为梯度张量在不调用 model.zero_grad() 或 optimizer.zero_grad() 的情况下不会重置。如果损失在训练样本上要取平均，我们还需要除以累积步骤的数量。

以下是使用梯度累积训练模型的要点。在这个例子中，我们可以用一个大于 GPU 最大容量的 accumulation_steps 批量进行训练：

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

2 充分利用多 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

3 多 GPU 机器上的均衡负载

DataParallel 有一个问题：GPU 使用不均衡。即在一些设置下，GPU-1 会比其他 GPU 使用率高得多。
使用 torch.nn.DataParallel 的前向和后向传播。
在这里插入图片描述
在前向传播的第四步（右上），所有并行计算的结果都聚集在 GPU-1 上。这对很多分类问题来说是件好事，但如果你在大批量上训练语言模型时，这就会成为问题。

我们可以快速计算语言模型输出的大小：

语言模型输出中的元素数量。

假设我们的数据集有 4 万词汇，每一条序列有 250 个 token、每个 batch 中有 32 条序列，那么序列中的每一个元素需要 4 个字节的内存空间，模型的输出大概为 1.2GB。要储存相关的梯度张量，我们就需要把这个内存翻倍，因此我们的模型输出需要 2.4GB 的内存。

这是典型 10GB GPU 内存的主要部分，意味着相对于其它 GPU，GPU - 1 会被过度使用，从而限制了并行化的效果。

如果不调整模型或优化方案，我们就无法轻易减少输出中的元素数量。但我们可以确保内存负载在 GPU 中更均匀地分布。

解决办法是把每部分输出保留在其 GPU 上，而不是将它们聚集到 GPU-1 上。我们也需要分配损失标准计算，计算损失并进行反向传播。

幸而，张航开源了一个名为 PyTorch-Encoding 的 PyTorch 包，它包含了这些定制的并行化功能。

我提取并稍稍改动了这个模块，你可以从以下地址下载 gist（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_gpu 张量的元组，每个张量分布在相应的 GPU 上。

DataParallelCriterion 容器封装了损失函数，并把 n_gpu 张量元组和目标标签张量作为输入。它在每个 GPU 上并行计算损失函数，像 DataParallel 分割模型输入一样分割目标标签张量。

下图说明了 DataParallelModel/DataParallelCriterion 的内部情况：
在这里插入图片描述
使用 DataParallelModel 和 DataParallelCriterion。

以下是你可能会遇到的两个特定案例的解决办法：

你的模型输出几个张量：你可能想分解它们：output_1, output_2 = zip(*predictions)

有时候你并不想使用并行损失函数：收集 CPU 上的所有张量：gathered_predictions = parallel.gather(predictions)

4 分布式训练：在多台机器上训练

在更大的批量上训练时，我们要如何控制多个服务器的算力呢？

最简单的选择是使用 PyTorch 的 DistributedDataParallel，它几乎可以说是以上讨论的 DataParallel 的直接替代元件。

但要注意：尽管代码看起来很相似，但在分布式设定中训练模型要改变工作流程，因为你必须在每个节点上启动一个独立的 Python 训练脚本。正如我们将看到的，一旦启动，这些训练脚本可以通过使用 PyTorch 分布式后端一起同步化。

在实践中，这意味着每个训练脚本将拥有：

它自己的优化器，并在每次迭代中执行一个完整的优化步骤，不需要进行参数传播（DataParallel 中的步骤 2）；

一个独立的 Python 解释器：这也将避免 GIL-freeze，这是在单个 Python 解释器上驱动多个并行执行线程时会出现的问题。

当多个并行前向调用由单个解释器驱动时，在前向传播中大量使用 Python 循环/调用的模型可能会被 Python 解释器的 GIL 放慢速度。通过这种设置，DistributedDataParallel 甚至在单台机器设置中也能很方便地替代 DataParallel。

现在我们直接讨论代码和用途。

DistributedDataParallel 是建立在 torch.distributed 包之上的，这个包可以为同步分布式运算提供低级原语，并能以不同的性能使用多种后端（tcp、gloo、mpi、nccl）。在这篇文章中，我将选择一种简单的开箱即用的方式来使用它，但你应该阅读文档和 Séb Arnold 写的教程来深入理解这个模块。

文档：https://pytorch.org/docs/stable/distributed.html
教程：https://pytorch.org/tutorials/intermediate/dist_tuto.html

我们将考虑使用具有两个 4 - GPU 服务器（节点）的简单但通用的设置：
在这里插入图片描述
主服务器（服务器 1）拥有一个可访问的 IP 地址和一个用于通信的开放端口。

改写 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 地址（我们的案例中是 192.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)                                     # Optimizer step

在第二台机器上，我们类似地启动脚本：

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 参数，其在第一台机器上被设为 0，在第二台机器上被设为 1（如果再加一台机器，则设为 2，以此类推…）。

资料参考：
1 GPU捉襟见肘还想训练大批量模型？谁说不可以
2 Pytorch多GPU并行处理

mingo_敏

关注

4
点赞
踩
27

收藏

觉得还不错? 一键收藏
4
评论
基于PyTorch使用大batch训练神经网络

如果你的 GPU 只能处理很少的样本，你要如何训练大批量模型？1 在一个或多个 GPU 上训练大批量模型当增大batch size时，经常会遇到CUDA RuntimeError：内存不足，但是为了提高性能，我们又需要增大batch size时，可以试试梯度累积。梯度下降优化算法的五个步骤。与之对等的 PyTorch 代码也可以写成以下五行：predictions = model...
复制链接

扫一扫