高度优化,京东AI开源的二值网络inference框架

导语:京东 AI 开源了一个针对 arm 指令集高度优化的二值网络 inference 框架 dabnn。据介绍,这是第一个高度优化的针对二值网络的开源推理框架,和 BMXNet 相比,dabnn 的速度有一个数量级的提升,对应的论文已被多媒体顶会 ACM MM 的 Open Source Software Competition 接收。

二值网络(Binary Neural Networks,又称 1-bit 网络)是一种特殊的神经网络,它将网络的权重和中间特征压缩为 1 个 bit,可以看作是普通浮点型网络量化到极致的结果。和其它 bit 数稍高的量化网络(例如三值网络、2-bit 网络、4-bit 网络)相比,二值网络最突出的优点在于,1-bit 乘加操作可以通过位运算高效实现,因此可以无缝运行在主流硬件平台(x86、ARM)上,而三值、2-bit、4-bit 等等操作需要特殊硬件平台(FPGA、ASIC)支持。在计算效率上,一个 32 位浮点运算指令只能完成一次浮点乘法或加法运算,但一个 32 位位运算指令可以同时完成 32 个 1-bit 乘法运算,具体来说,一个二值卷积可以通过 xnor(异或) 和 bitcount(计算操作数中值为 1 的 bit 的数量)操作高效完成:

在二值网络领域已经有一些推断框架,例如 BMXNet [2],BitStream [3],BitFlow [4]。其中,BitStream 和 BitFlow 只有论文,而没有源代码或二进制库,BMXNet 虽然开源,但在 Google Pixel 1 手机上进行的测试显示,它甚至比全精度推断框架 TensorFlow Lite 还要慢。

dabnn

为了填补这个空白,京东 AI 开源了一个针对 arm 指令集高度优化的二值网络 inference 框架 dabnn,这是第一个高度优化的针对二值网络的开源推理框架,和 BMXNet 相比,dabnn 的速度有一个数量级的提升,对应的论文已被多媒体顶会 ACM MM 的 Open Source Software Competition 接收。

论文链接:https://arxiv.org/abs/1908.05858

项目地址:https://github.com/JDAI-CV/dabnn

Demo链接:https://github.com/JDAI-CV/dabnn-example

dabnn 能取得非常快的速度的原因之一是,dabnn 的作者发现,虽然GEMM 是实现卷积的通用方法,BMXNet 和 BitStream 也用 GEMM 的 1-bit 变种 —— BGEMM 实现二值卷积,但对于 ARM 架构的二值卷积来说,它并不是最优的方法。因为 arm 架构下的 bitcount 指令有着特殊的特性,每个 bitcount 之后需要添加一个称作 addv 的额外指令,才能得到操作数中值为 1 的 bit 数量,这将导致计算二值卷积所需的指令数大大增加。基于这个观察,dabnn使用Binary Direct Convolution而不是 BGEMM实现二值卷积,经过仔细的设计,addv 的数量被减少到了最小,降低了推断的延迟。

为了证明 Binary Direct Convolution 的有效性,dabnn 的作者将 BGEMM 中所有的 addv 指令去掉,得到一个非正常的 BGEMM 实现,并比较了 BGEMM、BGEMM without addv 和 Binary Direct Convolution 的速度。对比结果清晰的说明过多的 addv 指令大大拖慢了 BGEMM 的速度(如图所示)。

图表 1 不同的二值卷积方法的比较(图片来自于 [1])

和已有框架的比较

如前文所说,在二值网络领域,几个已有 inference 框架或是不开放使用,或是速度较慢。与之相比,dabnn 不仅开源,并且有非常快的推理速度。从下图的测试数据可以看出,针对 3x3 卷积,和 TensorFlow Lite(全精度)相比,dabnn 要快 8~10 倍,和 BMXNet (二值)相比,dabnn 要快 7~23 倍。图表 2不同的推断框架的速度比较,二值网络(如 Bi-Real Net)通常含有一部分全精度卷积,因此加速比没有单个二值卷积那么明显(图片来自于[1])

图表 3各个二值网络推断框架的比较(图片来自于[1])如何转换模型

dabnn 同时开源了转换 ONNX 模型为 dabnn 模型的转换工具(值得一提的是,dabnn 的作者同时也是 ONNX 的资深贡献者),ONNX 是一个被几乎所有训练框架支持的模型格式。作为对比,BMXNet 重度依赖于 MXNet 且只支持 MXNet 模型。

使用 dabnn 的项目

dabnn 的 paper 放到 arxiv 不到两个月,已经有两篇二值网络 paper 使用 dabnn 进行实验,分别是商汤的 IR-Net: Forward and Backward Information Retention for Highly Accurate Binary Neural Networks(https://arxiv.org/abs/1909.10788 )和北航等机构的 Balanced Binary Neural Networks with Gated Residual (https://arxiv.org/abs/1909.12117 )。这体现了 dabnn 对于二值网络领域的重要性。

产业框架inference二值网络京东
暂无评论
暂无评论~