tvm
首发于tvm

量化深度学习模型在CUDA平台的自动优化

深度学习已经成功地在各种任务中得到应用。模型的推断速度在无人驾驶等实时的场景下尤为关键。网络量化是加速深度学习模型一种有效的方法。在量化的模型中,我们使用 int8、float16 等低精度的数据类型表示数据和模型参数。更低的数据带宽不仅降低了模型推断时间和内存及存储上的需求,也降低了运行时的功耗。与此同时,通过采用合适的 calibration (校准)方案,我们能够把模型精度损失降到最低。模型量化使得较大的模型能够在GPU、CPU、移动端设备等不同的硬件上部署。因此,模型量化在近年来受到大量的研究者和开发者的关注。

在此之前,量化模型中的运算符通常使用人工编写的 microkernel 来实现,或者依赖cuDNN、TensorRT等以黑盒形式提供的商用软件。通过汇编编写高性能的 microkernel 是一项困难且耗费大量人力的工作;并且,这些 microkernel 通常无法在新出现的模型和设备上达到较好的性能。

图1: 不同模型在TVM、TensorRT、MXNet上的运行时间

针对这个挑战,TVM 使用一个完整的编译器栈和基于机器学习的优化器来自动地生成高效的代码。通过在人工设计的搜索空间内自动搜索,TVM 能够生成高效的 kernel 代码。在VGG、ResNet等标准的模型上,TVM 达到了和 TensorRT 近似的性能。而在ResNext、Deformable ConvNets等新出现的模型上,利用TVM的自动优化我们能够十分容易地实现其他框架无法做到的性能提升。

我们将在本篇文章中介绍 TVM 在 CUDA 平台上对 int8 量化深度学习模型的优化。

使用 TVM 实现量化 CUDA kernel

通过 tensorization 利用 intrinsic 指令

许多平台提供了特定架构下的优化指令来处理一些特殊的计算。例如,x86 平台的 SIMD 指令、CUDA 平台的 dp4a 和 hfma 指令等。这些指令通常高度优化。利用这些 intrinsic 指令,我们能够实现显著的性能提升。

我们在 TVM 的 CUDA int8 运算符中使用 dp4a 指令进行优化。dp4a 是 CUDA 平台 Compute Capability 6.1 中的用于计算长度为 4 的 8 位整数向量点积,并产生32位计算结果的一条 intrinsic 指令。我们使用 dp4a 实现任意满足长度为 4 的倍数的 int8 向量点积。而卷积、全连接等通常依赖点积运算,借助高效的向量点积我们能够更进一步地实现这些更高层次的运算符。

例如,2D卷积可以直接在 channel、width、height 维度累加进行计算。这是一个 dp4a 的典型应用。在TVM中,我们使用 tensorization(张量化)来支持外部指令的调用。这一步不需要修改计算过程的声明,只需在 TVM 中的 schedule 阶段调用 “tensorize” 这一个 schedule primitive,将累加部分替换为基于 dp4a 的实现。关于 tensorization 的细节可以查看这篇教程

数据布局的重排

在做 tensorization的时候,我们可能需要对计算逻辑进行特殊的设计。在全连接运算符中,我们可以直接沿着数组内层索引的顺序依次累加。而卷积运算的实现更为复杂。在2d卷积中,我们希望沿着 channel 这一维度使用 dp4a 指令进行运算,因为大多数情形 channel 是 4的倍数 (否则将使用 NCHW 布局下的默认实现)。同时,为了更好的内存局部性,我们希望沿着最内层的轴进行累加。考虑到这些因素,我们使用一种自定义的布局。

在 CUDA int8 conv2d 中,更加经验我们选用了 NCHW4c 作为数据布局,OIHW4o4i 作为卷积核的布局。实际上,我们的实现能够很容易地推广到 NCHW[x]c 和 OIHW[x]o[x]i 这种更一般的情况。其中,x 是任意可被 4 整除的正整数。在我们选用的 NCHW4c 数据布局中,我们沿 channel 维度将每 4 个元素打包放到最内层。类似地,我们将卷积核中,我们将对应输入和输出 channel 的维度也分别进行打包并移到最内层。这种卷积核布局使得输出具有和输入相同的布局。因此,我们不需要在每层之间插入不必要的布局转换。

图2 展示了计算 2d卷积输出的一个元素的过程。在我们选用的布局(NCHW4c和OIHW4o4i)中,按 NCHW 或 OIHW 四个维度索引得到的结果是打包后的元素。图2中间,卷积核部分的每一列来自于原来卷积核中不同的filter。我们用 dp4a 计算中间这幅图packed kernel部分每一行与 packed input 的点积并累加到输出结果上。

图 2. 在 NCHW4c、OIHW4o4i 布局下的 2D 卷积。左:NCHW4c 布局下的输入。蓝色部分为一个移动的 filter。灰色部分对应输入及卷积核的一个元素。中: 左图灰色部分的一个元素实际对应输入数据和卷积核中多个元素。右: NCHW4c 布局下的输出。图中画出的一个元素,对应 channel 维4个元素。

在指定卷积层的布局之后,在 Relay 的 AlterOpLayout pass中,我们会自动地将其他的运算符如张量加法、激活函数等转换到所选的布局。对卷积核的布局转换可以离线地提前计算。我们可以让整个网络在新的数据布局下运行,而没有额外的开销。

设计自动优化的搜索空间

我们所实现的量化版本的运算符达到好的性能的关键是使用基于机器学习的自动优化。在这里,我们需要考虑一个问题:如何设计一个有效地 schedule 的搜索空间?一个有效的 schedule 搜索空间是指自动调优过程能够在合理的的迭代次数下达到理想的性能。通常来说,我们需要尽量地让搜索所使用的模版更加灵活,尽可能地覆盖搜索空间中不同的配置。另一方面,我们也需要借助性能优化上的先验知识。例如,把数据缓存到 shared memory 是 CUDA 编程的一个常用技巧。因此,我们在使用 shared memory 的同时,让机器学习算法来自动的选择 shared memory 大小的最优分配。我们也手动地做了一些 tiling 例如把一些维度按 4 或 16的因子进行分割来以利用向量化的内存访问。

在 int8 2d 卷积算法中,我们设计的搜索空间包含了一组可调的选项,例如 tile size,融合多层循环时的选择,loop unrolling (循环展开) 及 double buffering (双缓存)的设置。TVM 中,CUDA平台量化版本的 conv2d 和 dense 注册在 int8 这一 template key下。在做自动调优的时候,我们可以在创建任务的时候设置 template key 来选择这一实现。关于自动调优的细节可以查看这篇AutoTVM教程.

总体工作流程

图 3. 运行量化模型的流程

TVM 提供了简单的工作流程。只需几步,即可量化从其他框架训练好的模型、自动优化模型中的运算符(基于AutoTVM),并部署到不同的设备上。

首先,我们使用 Relay 前端导入现有的模型。在此我们使用一个输入形状为 (1, 3, 224, 224) 的MXNet模型为例。

sym, arg_params, aux_params = mxnet.model.load_checkpoint(model_path, epoch)
net, params = relay.from_mxnet(sym, shape={'data': (1, 3, 224, 224)}, arg_params=arg_params, aux_params=aux_params)

接下来,我们直接调用 Relay 的量化 API 生成量化后的模型。

net = relay.quantize.quantize(net, params=params)

得到量化模型之后,我们使用 AutoTVM 对模型中使用的运算符创建调优任务并进行自动优化。该部分可以参考AutoTVM教程

最后,我们编译最终得到的模型,即可在量化模式下运行。

with relay.build_config(opt_level=3):
    graph, lib, params = relay.build(net, target)

relay.build 的返回值可以直接用于部署。我们可以直接在本地GPU上运行,或通过 RPC 部署到远端设备.

性能测试

为了验证 TVM 中量化版本的运算符的性能,我们在一些常用的模型如 VGG-19、ResNet-50、Inception V3 上进行测试。此外,我们还测试了DRN-C-26, ResNeXt-50, 和Deformable ConvNets中的 DCN-ResNet-101 来比较新出现的模型上的性能。这些模型包含一些较少使用的运算,如 dilated 卷积、 group 卷积、deformable 卷积等。我们选用 NVIDIA TensorRT 作为 baseline. 提供的 MXNet 1.4 + cuDNN 7.3 在 float32 模式下的数据可以作为量化加速的参考。该测试是在一张 NVIDIA GTX 1080 上进行的。我们报告在 batch size = 1 和 16下,平均每张输入图片的运行时间。

如图1所示,相比于 float32 的模型,使用 TVM 量化后能够达到最高8倍的加速。在VGG、ResNet等标准的 CNN 模型上,TVM 达到了近似于 TensorRT 的 state-of-the-art 结果。

TVM 在新出现的模型上也达到了很好的结果。在 ResNeXt 和 DCN-ResNet-101 上 TVM 有显著的性能提升。因为没有 deformable 卷积的官方实现,我们暂无 TensorRT 在 DCN-ResNet-101 上的结果。可以看出,自动优化使得 TVM 能够简单且灵活地支持和优化新出现的模型。

相关代码

相关链接

编辑于 2019-05-02

文章被以下专栏收录

    深度学习编译优化和硬件部署

    上海交通大学校园内的计算机技术组织,专栏主要是来自团队成员的投稿,多为技术话题,偶尔伴有分享通告等内容。欢迎加入我们。