如何将 Tensorflow.js 的性能疯狂提升 100%+

编者按：本文由支付宝体验技术部数据智能团队成员青壁编写。文中内容是他根据在实际项目中总结而来。

在正式开始前，我们先来看看两组对比的数据。

模型首次执行性能对比  tfjs vs ant-tfjs

推理性能对比（基于MobileNetV2）

1. Web上的高性能计算

[R]. Web Worker

使用 Web Worker 可以将一些 CPU 密集型计算转移到子线程中去做执行，同理可以将计算进行分拆，创建多个线程进行并行计算。这里为各位大佬献上一个利用 Web Worker 做并行计算的库 Paralles.js。

但是即使是支持了多线程并行计算，由于 JS 动态语言的特性，即使有了 V8引擎的加持，也是远远满足不了我们执行深度学习模型的性能需要。

[G]. Asm.js

2012年，Mozilla 的工程师 Alon Zakai 在研究 LLVM 编译器时突发奇想：许多 3D 游戏都是用 C / C++ 语言写的，如果能将 C / C++ 语言编译成 JavaScript 代码，它们不就能在浏览器里运行了吗？众所周知，JavaScript 的基本语法与 C 语言高度相似。

于是，他开始研究怎么才能实现这个目标，为此专门做了一个编译器项目 Emscripten。这个编译器可以将 C / C++ 代码编译成 JS 代码，但不是普通的 JS，而是一种叫做 asm.js 的 JavaScript 变体。

C / C++ 编译成 JS 有两个最大的困难。

• C / C++ 是静态类型语言，而 JS 是动态类型语言。
  
• C / C++ 是手动内存管理，而 JS 依靠垃圾回收机制。
  

asm.js 就是为了解决这两个问题而设计的：它的变量一律都是静态类型，并且取消垃圾回收机制。除了这两点，它与 JavaScript 并无差异，也就是说，asm.js 是 JavaScript 的一个严格的子集，只能使用后者的一部分语法。

一旦 JavaScript 引擎发现运行的是 asm.js，就知道这是经过优化的代码，可以跳过语法分析这一步，直接转成汇编语言。另外，浏览器还会调用 WebGL 通过 GPU 执行 asm.js，即 asm.js 的执行引擎与普通的 JavaScript 脚本不同。这些都是 asm.js 运行较快的原因。据称，asm.js 在浏览器里的运行速度，大约是原生代码的50%左右。

Asm.js 相关介绍摘自阮一峰老师的《asm.js 和 Emscripten 入门教程》

即使如此，asm.js 依旧满足不了执行深度学习模型的性能需要。

[B]. WebAssembly

关于 Wasm 的文章现在实在是太多了，我这里也就不赘述了。

Wasm 相对于原生 JS 或者 asm.js 来说，速度的确很快，TensorFlow.js 官方目前也实现了基于 Wasm 的 backend。但是目前在绝大多数的机器上，Wasm 的执行速度相较于下面要说到的 WebGL 要差个3倍以上。尤其在一些大型模型上表现更为明显。

上图是当前的 Wasm 和 WebGL 的性能对比，可以看到，WebGL 比 Wasm 要快5~7 倍左右。Wasm 最根本的问题还是在 width 较大的模型上，并行度和 GPU 相比还是不够，这个是 CPU 和 GPU 的硬件结构决定的，所以说在较大的模型上 WebGL 或者未来的 WebGPU 依旧是首选的 backend。

不过随着 Wasm 对 SIMD 指令集的支持以及多线程的支持，以后在性能上与 WebGL 的差距会越来越小。

[A]. GPU

在浏览器内利用 GPU 的能力有两种方法：

1. 十分成熟的 WebGL
   
2. 尚在草案阶段的 WebGPU。当前 WebGPU 的进展十分缓慢，我们这里先抛开不谈，谈一谈如何利用 WebGL 进行前端的高性能计算。
   

说起 WebGL，大家的第一反应应该就是“这不是用来做图形渲染的吗？和高性能计算有啥关系？”。各位看官莫急，且往下看。

在 Web 上面利用 WebGL 做高性能计算其实是依赖于一个 offscreen 的 canvas。canvas 有许多个像素组成，每个像素的颜色可以有 RGBA 四个维度表示，每个维度范围为0-255既8位。把 RGBA 表示成数值的话，那每个像素可以存32位。这就是前端使用 GPU 计算最为核心的一点，每个像素可以存储一个32位的值, 刚刚好就是一个 int 或者 uint。

首先我们看下 WebGL 的渲染流程：

其中两个 vertex shader 和 fragment shader 为两个 GLSL 代码片段，分别处理坐标数据和颜色数据。vertex shader 和 fragment shader 的执行是以像素为单位，canvas 开始绘制的时候 vertex shader 中得到。每个需要绘制的像素的坐标，视需要可以对坐标进行各种转换，最终得到一个最终位置。这个过程中可以将数据作为输出传 fragment shader 参与下一步的计算 fragment shader 接受各种输入，最终输出一个 RGBA 颜色数据作为该像素点的颜色值。

当所有像素都绘制完成之后，画布绘制完成。

tfjs 就是利用这种方式实现了 Web 内执行模型推理的加速。

 

2. ant-tfjs的WebGL的优化

接下来从两个方面介绍下我们的优化方案：首次执行(预热)性能、推理性能。

首次执行(预热)优化

模型的首次执行是指模型第一次在手机上执行的时间，原因是由于目前的 tfjs 计算使用 GPU(WebGL) 进行计算加速。(WebAssembly 目前性能不够，所以 WebGL 是目前唯一可用的方案)。

熟悉 WebGL 开发的朋友应该都清楚，WebGL 的计算主要依赖于 Shader(着色器)，一般一个模型中的计算节点(op)至少依赖一个 FragmentShader。而对于一个中等 MobilenetV2 的模型来说，一般会存在大几十甚至100+个计算节点，那么我们在首次运行的时候需要对 Shader 进行编译、执行、寻址、缓存等等各种操作，所以会导致首次执行性能很慢。而由于首次执行时更是涉及到了模型文件的各种编码(在 GPU 内完成)，所以导致 Shader 数量巨大。

而从图上可以看到，ant-tfjs 相对于官方的 tfjs 在首次执行性能上有了巨大的提升，80%~100%+。具体是怎么做到的呢？且听我娓娓道来。

 

一个模型的执行其实就是一张有向无环图的顺序执行，每次的执行就是图中的一个计算节点。下面以一个典型的计算节点（卷积）举例：

计算节点可以表示为：Y = Conv(X, W)。Conv 为卷积算法，接受 X 和 W 两个输入(实际可能还会有bias、preluActivation等变量输入)。X 是上一个节点(xxx OP)的输出，在 GPU 内表现为一块已编码的 texture，而 W 则是从我们模型文件中直接读取的数据。

 

由于计算要在 GPU 内进行，所以需要把模型文件内的数据上传至 GPU 内(通过texture存储)，由于 texture 的数据存储的特殊性，所以需要对原始模型数据进行编码，使之能够适应 GPU 的计算要求。所以在首次执行时，tfjs 会对每一个节点的W(权重)数据进行上传、编码操作，这样就产生了大量的 Shader。

所以，我们只要对模型文件进行预先编码，在离线环境对模型数据进行排布，这样就能省去了 GPU 内数据编码这一步，从而获得了巨大的性能收益。

 

推理优化

影响推理性能的因素事实上非常多，从 WebGL 渲染管线到 JS 执行，方方面面都会拖慢性能，这里先举上几个🌰：

• 计算节点过多，导致频繁切换 WebGL program；
  
• 数据在 GPU 内的内存布局不合理，导致 L1 cache 频繁 miss；
  
• 分支过多，影响 GPU 计算单元并行化；
  
• 对 GPU 并行化特性未充分利用;
  
• JS 代码未针对 jitless 做优化，导致 iOS 小程序环境下 JS 代码执行过慢；
  
• ...
  

 

问题分析清楚了，那么再去解决就事半功倍了。

 

图优化

图优化手段还是挺多的，具体可以参考 TVM，我们这里实际上是用到了 OP FUSION。将一些可以融合的节点在图结构上进行融合(nOP -> 1OP)，基于新的计算结点实现新的 OP。这样一来大大减少了 OP 的数量，进而减少了 Program 的数量，所以提升了推理性能。在低端手机上效果尤为明显。

向量化

GPU 计算能力强大的原因就在于它的访存带宽及计算并行化，官方的 tfjs 在并行化计算上做的并不足，在一些高频率、高计算量的 OP 上仍然采用的是逐点计算的方法。所以针对这些 OP 进行向量化优化，相对充分利用 GPU 的并行化能力，对提升推理性能的优化起到很大的作用。

 

jitless 优化

jitless 优化主要针对于 iOS 小程序无 JIT 的场景，就是一些常规的 JS 性能优化方案，这里不做赘述了。

 

优化内存布局

相对于并行计算的优化，访存的效率更容易成为推理性能的瓶颈。相对一个内存布局糟糕的方案，一个优秀的内存布局性能可以是其数倍。比如各种推理引擎常见的 IM2COL 算法，就是一种优化卷积操作访存连续性的算法。官方 tfjs 的内存布局是对一个矩阵中的元素进行2x2的 pack，即将相邻的两个 row 和相邻的两个 colpack到一个像素点内。这种布局方法实际上并未充分利用 GPU 的 cache 机制，尤其在矩阵乘法这种 high traffic 的场景下，会导致 cache 频繁 miss，这样就带来了糟糕的性能。

 

所以，通过以上的种种操作，我们最终将 ant-tfjs 的性能(预热、推理)相对于官方 tfjs 最高提升了100%+。并且在宠物宝项目上取得了很好的效果。

 

更多的优化方案

目前的计算虽然是基本已经向量化了，但是依旧有一定的优化空间。例如：

• 提升每次渲染时的 texture bandwidth，继而提升访存及并行效率；
  
• 并行渲染。
  

我们优化还在继续，我们有信息能够将性能再提升一个数量级。希望未来在低端手机上也能看到 30fps 的 mobilenet。

如果你对前端智能化或者 AntTF.js 感兴趣，想参与 AntTF.js 的开发，欢迎发送邮件至：diforce-talent@antgroup.com。