模型训练太慢？显存不够？这个方法让你的GPU联手CPU

随着深度学习模型复杂度和数据集规模的增大，计算效率成为了不可忽视的问题。GPU凭借强大的并行计算能力，成为深度学习加速的标配。然而，由于服务器的显存非常有限，随着训练样本越来越大，显存连一个样本都容不下的现象频频发生。除了升级硬件（烧钱）、使用分布式训练（费力），你知道还有哪些方法吗？即使显存充足，所有运算都在GPU上执行就是最高效吗？只要掌握以下小知识，模型训练的种种问题统统搞定，省时省力省钱，重点是高效！

其实CPU和GPU是协同工作的，如果能合理地利用它们各自的优势，就能够节省显存资源（显存不够内存来凑），甚至获得更好的训练性能。本文为您提供了device_guard接口，只需要一行命令，即可实现GPU和CPU的混合训练，不仅可以解决训练模型时通过调整批尺寸（batch size）显存依然超出的问题，让原本无法在单台服务器执行的模型可以训练，同时本文还给出了提高GPU和CPU混合训练效率的方法，将服务器资源利用到极致，帮助您提升模型的性能！

模型训练的特点

深度学习任务通常使用GPU进行模型训练。这是因为GPU相对于CPU具有更多的算术逻辑单元（ALU），可以发挥并行计算的优势，特别适合计算密集型任务，可以更高效地完成深度学习模型的训练。GPU模式下的模型训练如图1所示，总体可以分为4步：

第1步，将输入数据从系统内存拷贝到显存。

第2步，CPU指示GPU处理数据。

第3步，GPU并行地完成一系列的计算。

第4步，将计算结果从显存拷贝到内存。

图1 模型训练示意图

从图中可以了解到，虽然GPU并行计算能力优异，但无法单独工作，必须由CPU进行控制调用；而且显存和内存之间的频繁数据拷贝也可能会带来较大的性能开销。CPU虽然计算能力不如GPU，但可以独立工作，可以直接访问内存数据完成计算。因此，想获得更好的训练性能，需要合理利用GPU和CPU的优势。

 

模型训练的常见问题

 

问题一：GPU显存爆满，资源不足

 

你建的模型不错，在这个简洁的任务中可能成为新的SOTA，但每次尝试批量处理更多样本时，你都会得到一个CUDA RuntimeError：out of memory。

这是因为GPU卡的显存是非常有限的，一般远低于系统内存。以V100为例，其显存最高也仅有32G，甚至有些显存仅12G左右。因此当模型的参数量较大时，在GPU模式下模型可能无法训练起来。

设置CPU模式进行模型训练，可以避免显存不足的问题，但是训练速度往往太慢。

那么有没有一种方法，可以在单机训练中充分地利用GPU和CPU资源，让部分层在CPU执行，部分层在GPU执行呢？

 

问题二：频繁数据拷贝，训练效率低

 

在显存足够的情况下，我们可以直接采用GPU模式去训练模型，但是让所有的网络层都运行在GPU上就一定最高效吗？其实GPU只对特定任务更快，而CPU擅长各种复杂的逻辑运算。框架中有一些OP会默认在CPU上执行，或者有一些OP的输出会被存储在CPU上，因为这些输出往往需要在CPU上访问。这就会导致训练过程中，CPU和GPU之间存在数据拷贝。

图2是CPU和GPU数据传输示意图。假设模型的中间层存在下图中的4个算子。其中算子A和算子B都在CPU执行，因此B可以直接使用A的输出。算子C和算子D都在GPU上执行，那么算子D也可以直接使用C的输出。但是算子B执行完，其输出在CPU上，在算子C执行时，就会将B的输出从CPU拷贝到GPU。

频繁的数据拷贝，也会影响模型的整体性能。如果能把算子A和B设置在GPU上执行，或者算子C和D设置在CPU上执行，避免数据传输，或许会提升模型性能。那么应该如何更加合理地为算子分配设备，使得训练过程更加高效呢？我们需要更综合地考虑，在发挥GPU和CPU各自计算优势的前提下，降低数据拷贝带来的时间消耗。

图2 CPU和GPU数据传输示意图

 

device_guard自定义GPU和CPU混合训练

 

上面两个场景都是希望为模型中的部分层指定运行设备。飞桨提供了fluid.CUDAPlace和fluid.CPUPlace用于指定运行设备，但这两个接口在指定设备时是二选一的，也就是说要么在GPU模式下训练，要么在CPU模式下训练。过去我们无法指定某一部分计算在GPU上执行还是在CPU上执行。飞桨开源框架从1.8版本开始提供了device_guard接口，使用该接口可以为网络中的计算层指定设备为CPU或者GPU，实现更灵活的异构计算调度。

如何使用device_guard接口解决上面两个场景中提到的问题呢？接下来，我们看看具体的例子。

 

好处一：充分利用CPU资源，避免显存超出

 

如果使用fluid.CUDAPlace指定了全局的运行设备，飞桨将会自动把支持GPU计算的OP分配在GPU上执行，然而当模型参数量过大并且显存有限时，很可能会遇到显存超出的情况。如下面的示例代码，embedding层的参数size包含两个元素，第一个元素为vocab_size(词表大小)，第二