使用 TVM 将深度学习模型编译到 WebGL

•


现在 TVM 推出一个全新的 OpenGL/WebGL 后端！这篇博文将解释它是什么，以及您可以使用它实现什么。

OpenGL/WebGL 后端

TVM 已经针对许多后端，涵盖各种平台：CPU、GPU、移动设备等…… 这次我们添加了另一个后端：OpenGL/WebGL。

OpenGL/WebGL 使我们能够在没有安装 CUDA 的环境中使用 GPU。在目前看来，这是在浏览器内部使用 GPU 的唯一方法。

这个新的后端允许我们以 3 种不同的方式使用 OpenGL/WebGL

  • 本地 OpenGL：我们可以将深度学习模型编译成 OpenGL，并直接在本地机器上运行它，完全使用 Python。
  • WebGL 与 RPC：我们可以将深度学习模型编译成 WebGL，并通过 Emscripten 将其导出为共享库，包括 JavaScript 主机代码和 WebGL 设备代码。然后，我们可以通过 RPC 将该库部署到浏览器内部运行的 TVM JavaScript 运行时系统上。
  • WebGL 与静态库：我们可以将深度学习模型编译成 WebGL，将其与 TVM JavaScript 运行时系统链接，并导出整个软件包。然后，我们可以在浏览器中的网页上运行该模型，而无需任何依赖项。详细流程如图 1 所示。

我们依赖 Emscripten 及其 fastcomp LLVM 后端来生成 javascript 后端。

image
图 1

请参阅此处，查看所有这三种方式的示例。

这与 X 有何不同？

在浏览器上运行神经网络并不是一个全新的事物。Andrej Karpathy 的 ConvNetJS 和 Google 的 DeepLearning.JS 就是这样的例子。

那么 TVM 与 WebGL 的独特之处是什么？最大的不同在于 TVM 中的算子内核是自动编译的，而不是手写的。如图 2 所示，TVM 利用统一的 AST 来定义内核，并将其编译为不同平台上的代码。


图 2

这意味着

  • 要将您现有的模型部署到 WebGL，您不需要编写大量额外的代码。NNVM/TVM 模型定义对于所有目标都是相同的，因此您只需要将其编译为新的目标即可。
  • 要添加新的算子内核，您只需要在 TVM 中定义一次，而不是为每个目标实现一次。您不需要知道如何编写 GLSL 代码来为 WebGL 添加新的算子内核！

基准测试

在这里，我们对一个典型的工作负载进行基准测试：使用 resnet18 进行图像分类。

我使用的是我的5 年旧的笔记本电脑，它配备了 8 核 Intel® Core™ i7-3610QM 和 GTX650M。

在这个基准测试中，我们从 Gluon 模型动物园下载了一个 resnet18 模型，并对一张猫的图像执行端到端分类。我们只测量模型执行时间（不包括模型/输入/参数加载），并且每个模型运行 100 次以获得平均值。结果如图 3 所示。

image
图 3

基准测试在 4 种不同的设置下运行

  • CPU (LLVM)：模型被编译成 LLVM IR 并进行 JIT 编译。因此，它完全在 CPU 上运行。
  • OpenCL：模型被编译成 OpenCL。仍然有一些胶水代码被编译成 LLVM，它负责设置和启动 OpenCL 内核。然后我们在本地机器上运行它。
  • OpenGL：与 OpenCL 相同，但编译为 OpenGL。
  • WebGL：胶水代码被编译成 LLVM，并使用 Emscripten 的 Fastcomp LLVM 后端转换为 JavaScript。设备代码被编译成 WebGL。我们在 Firefox 中执行模型。

从上面的结果我们可以观察到，TVM OpenGL 后端具有与 OpenCL 相似的性能。更有趣的是，浏览器内部的 WebGL 版本并没有比桌面 OpenGL 慢很多。考虑到主机代码是 JavaScript，这非常令人惊讶。这可能是由于 Emscripten 生成了 asm.js，它在 Firefox 中实现了显著的优化。

这是朝着自动将深度学习模型编译到 Web 浏览器的第一步。随着我们将优化引入 TVM 堆栈，我们期望获得更多的性能改进。

给我看代码

致谢

我们感谢 Emscripten 的开发者提供了 fastcomp 工具链，并在开发过程中提供了帮助。