加速计算机视觉：我们如何将EfficientNet扩展到IPU-POD超级计算系统上

2021/12/17

加速计算机视觉：我们如何将EfficientNet扩展到IPU-POD超级计算系统上

作者:

例如，与标准卷积操作相比，EfficientNet中的深度卷积或分组卷积既有很强的表现力，又有很高的计算效率，但其算术强度（计算与数据移动的比率）很低。由于存储和处理器核之间有大量的数据传输，这些类型的操作在GPU上表现不佳。IPU的架构以存储为中心，因此整个模型及其激活可以留在芯片上，减轻了这种昂贵的数据移动。

此外，IPU的MIMD（多指令、多数据）范式支持模型和训练程序在多个维度上的细粒度并行化。这有助于在使用小型卷积和并行处理极少的数据样本的同时，实现高吞吐量。相比之下，具有SIMD（单指令、多数据）架构的GPU在利用并行性方面受到限制，迫使用户在其机器学习算法的设计与硬件上可实现的算法吞吐量之间做出权衡。

这种在IPU硬件上的加速为更多的创新者和人工智能从业者提供了机会，使他们能够从EfficientNet模型的高效率中受益，从而在真实世界中大规模地实施。

Graphcore在GitHub上的示例提供了包括EfficientNet系列模型在内的大量主流模型，任何人都可以直接在IPU上使用。在本练习中，我们分析的是EfficientNet-B4，它常被用于人工智能领域对EfficientNet性能的基准测试。

优化吞吐量

分布模型

有很多种方法可以在一组IPU上分布一个模型，包括数据并行复制和流水线模型并行。鉴于EfficientNet-B4的尺寸和IPU-POD的规模，数据和模型并行性的某种组合可实现最佳结果。Poplar软件栈可以轻松地在框架级别尝试不同的分布式设置，从而加快我们寻找配置各种类型并行性的最佳方式。

GitHub示例中EfficientNet-B4的默认配置是通过四个IPU将模型流水线化。拥有多个流水线阶段意味着需要花费更多时间来填充和排空流水线，从而降低了整体IPU利用率。这也意味着需要更加谨慎地平衡IPU之间的工作量。如果我们能够将流水线阶段的数量减少到两个，这些挑战就可以得到解决。这也使我们能够将副本数量翻倍。

为了在两个IPU间拟合EfficientNet-B4，我们采用了三种技术：

使用16位浮点运算和主权重
减少局部批尺寸
使用G16版本的EfficientNet，如Graphcore Research（Graphcore研究院）发布的《使EfficientNet更高效》博客和论文中所述

IPU天然支持16位和32位浮点数表示，允许用户灵活地调整应用程序中存在的各种张量的表示。在某些情况下，可能需要存储FP32主权重，其中在整个训练过程中以全精度来存储模型参数的副本。我们发现，对于EfficientNet，当使用FP16主权重并在权重更新上使用随机舍入时，性能仍然得以保持。

由于储存激活所需的存储随批尺寸的增加而增加，因此必须考虑如何在不影响训练的情况下减少激活的开销。激活再计算是一种我们在需要激活进行反向传递时重新计算激活的方法。它提供了一个计算和存储之间的简单权衡。该方法在Poplar中可用，并且可以通过流水线API在框架层面上轻松访问。使用EfficientNet-B4，我们可以通过激活再计算来拟合一个尺寸为3的局部批量。与四级流水线设置相比，这种配置允许更多的样本被并行处理，因为我们有两倍的副本来分配微小的批。在这项工作中，我们使用Group Norm（组规范），一种与批次无关的归一化方法。由于这个模型没有任何批次间的依赖性，我们可以拥有在存储中拟合的任何局部批尺寸，并简单地增加梯度累积（GA）计数，以达到我们所希望的任何全局批尺寸，允许我们在不影响模型的训练动态的情况下在存储中拟合我们的激活。

通过将卷积组的大小从1增加到16（随后降低扩展率以补偿FLOP和参数的增加），我们减少了EfficientNet中MBConv块的存储开销。Graphcore Research（Graphcore研究院）的论文《使EfficientNet更高效》对此进行了更详细的探讨，并在我们的博客中进行了总结。除了节省存储，该模型的这种变体还具有提高ImageNet验证准确性的额外好处。

这三种技术使我们能够在两个IPU上拟合模型，而不必将任何优化器状态卸载到流存储器上，从而允许以高吞吐量进行训练。

数据IO（输入输出）

在强大的人工智能加速器系统（如IPU-POD）上训练机器学习模型时，一个常见的瓶颈是向模型提供足够的数据，以便从主机上进行处理。PopRun是一个命令行实用程序，通过在多个程序实例上以分布式的方式启动应用程序，帮助减轻这一瓶颈。每个实例的IO由其相应的主机服务器管理，使我们能够将模型向上扩展到各种IPU-POD系统，而不受限于向模型提供数据的速度。

虽然ImageNet数据集最初是以INT-8表示的，但作为预处理的一部分，它通常在主机上被转换为更高精度的浮点数据类型。通过应用这种转换，一旦数据在IPU上，通信开销就会减少，因为输入数据可以以较低的精度流向IPU。这有助于提高向模型输入数据的速度，进一步提高吞吐量。

IPU内部和IPU之间的数据移动是根据批量同步并行（BSP）执行方案实现的。通过BSP，tile（IPU处理器核）在局部计算和与其他tile的数据交换之间交替进行，中间有一个同步步骤。

这种执行方案使Poplar能够有效地在数十万个tile上进行并行操作。对于IO受限的应用来说，所有tile上的原始BSP范式会扼杀性能。Poplar通过允许一部分tile以异步方式专门用于数据流而其余的计算tile在BSP范式下执行来解决这个问题。重叠的IO缓解了IO的瓶颈，同时仍然允许高度可扩展的执行方案。我们发现，将每个IPU的1472个tile中的32个分配给IO重叠，并预取多达三批数据，再加上上述方法，在所有系统规模上都取得了令人印象深刻的吞吐量。

扩展批尺寸

为了在IPU-POD₁₂₈和IPU-POD₂₅₆等大型系统中进行训练，使用大的全局批尺寸是有益的。这使我们能够在许多副本上并行处理数据样本，同时确保每个副本有足够的工作来保持效率。此外，为了摊薄减少各副本梯度的成本，我们在各副本之间进行通信和更新权重之前，在局部积累多个向前和向后的梯度，这就是所谓的梯度累积（GA）。因此，全局批尺寸是局部批尺寸、副本数量和梯度累积计数的乘积。

我们的基准线来自《使EfficientNet更高效》，使用了768的全局批尺寸。给定局部批尺寸为3，在IPU-POD256上每次训练迭代只产生两个局部批尺寸，充分利用流水线设置，并且不得不应对频繁且昂贵的权重更新。然而，仅仅增加全局批尺寸会导致泛化性能的下降。这种现象在许多机器学习应用中很常见，特别是计算机视觉。因此，我们寻求的全局批尺寸要足够大，以保持所有IPU-POD系统的高吞吐量，但又要足够小，以使我们达到良好的统计效率。

众所周知，优化器对大批尺寸训练的稳健性有很大影响。最初的EfficientNet论文使用了RMSProp优化器，在我们的调查中，该优化器在批尺寸增加时很难保持统计效率。Wongpanich等人提出了类似的意见，并建议使用逐层自适应速率缩放（LARS），这种优化器在训练大批尺寸的视觉模型时效果良好。LARS在逐层的基础上对学习率进行缩放，以确保权重和权重更新之间有相似的幅度。除了LARS，我们还采用了通常与优化器一起使用的多项式衰减学习率计划。

通过对批尺寸、学习率、预热时期数、动量系数和权重衰减进行超参数扫描，我们发现，与最初的EfficientNet-B4实施相比，我们可以实现6144的全局批尺寸，而且性能没有任何损失。在这个批尺寸下，所有考虑到的IPU-POD都可以保持较高的梯度累积计数。

性能结果

通过灵活的超参数配置，我们现在能够以高吞吐量在一系列IPU-POD系统上训练EfficientNet-B4。根据最初的EfficientNet论文，我们训练了350个时期。尽管《使EfficientNet更高效》一文表明EfficientNet系列模型可以在以原始分辨率进行微调之前，以较低的图像分辨率进行预训练，但我们仍保留了最初的分辨率以保持与其他实施的可比性。