今天, 华盛顿大学陈天奇团队开发的 TVM 发布了更新, 不需要写任何 JavaScript 代码, 直接就能把深度学习模型编译到 webGL/OpenGL, 然后在浏览器运行
深度学习离不开 TensorFlow,MXNet,Caffe 和 PyTorch 这些可扩展深度学习系统, 但它们大多专门针对小范围的硬件平台 (例如服务器级 GPU) 进行优化, 要适应其他平台需要付出相当大的工程成本和费用, 这对深度学习系统的灵活部署提出了挑战
大量不同的深度学习框架(编程语言), 越来越多的硬件架构, 两者之间需要一个桥梁 TVM 框架正是为此而生, 旨在让研究人员和开发者能够在各种不同的硬件, 从手机嵌入式设备到低功耗专用芯片这些不同的系统上, 快速轻松地部署深度学习应用, 而且不会牺牲电池电量或速度
TVM 是神经网络和硬件后端之间一个共同的层(a common layer), 无需为每一类设备或服务器建立一个单独的基础架构, 该框架使开发人员能够在多种不同的硬件设备上快速轻松地部署和优化深度学习系统, 帮助研究人员快速优化新算法的实现, 验证新的思路, 还有助于促进深度学习研究领域的硬件和软件协同设计
新发布 OpenGL / WebGL 后端支持
TVM 已经支持多个硬件后端: CPU,GPU, 移动设备等...... 这次我们添加了另一个后端: OpenGL / WebGL
OpenGL / WebGL 使我们能够在没有安装 CUDA 的环境中利用 GPU 目前, 这是在浏览器中使用 GPU 的唯一方式
这个新的后端允许我们以一下 3 种方式使用 OpenGL / WebGL:
本地 OpenGL: 我们可以将深度学习模型编译成 OpenGL, 并直接在本地机器上运行, 完全只使用 Python
带有 RPC 的 WebGL: 我们可以将深度学习模型编译为 WebGL, 并将其作为一个共享库导出, 并带有 JavaScript 主机代码和 WebGL 设备代码然后, 我们可以通过 RPC 将这个共享库部署到 TVM JavaScript 运行时系统, 在浏览器内运行
带有静态库的 WebGL: 我们可以将深度学习模型编译为 WebGL, 将其与 TVM JavaScript 运行时系统连接, 并导出整个包然后, 可以在浏览器的网页中运行模型, 不需要依赖项详细流程如图 1 所示
图 1
以上三种方式的演示代码: https://github.com/dmlc/nnvm/blob/master/tutorials/from_mxnet_to_webgl.py
这与 X 有何不同?
在浏览器上运行神经网络并不是非常新奇的事 Andrej Karpathy 提出的 ConvNetJS, 以及 Google 的 DeepLearning.JS 都是这样的想法
那么使用 WebGL 的 TVM 有什么独特之处? 最大的区别是 TVM 中的操作内核是自动编译的, 而不是人工编译的如图 2 所示, TVM 使用统一的 AST 定义内核, 并将其编译为不同平台上的代码
图 2
这意味着:
你不需要编写大量附加代码, 就可以将现有模型部署到 WebGLNNVM / TVM 模型定义对于所有 target 都是相同的, 因此你只需将其编译到新的 target
如果要添加新的操作系统内核, 你只需要在 TVM 中定义一次, 而不用为每个 target 实现一次你不需要知道如何编写 GLSL 代码来向 WebGL 添加新的操作系统内核!
Benchmark
这里, 我们为一个典型的工作负载执行基准测试: 使用 resnet18 进行图像分类
我使用的是有 5 年历史的 8 核英特尔 ® 酷睿 i7-3610QM 笔记本电脑, 以及一个 GTX650M
在这个基准测试中, 我们从 Gluon 模型库里下载了 resnet18 模型, 并对猫的图像进行端到端分类我们只测量了模型执行时间(不包含模型 / 输入 / 参数的加载), 并且每个模型运行 100 次以取得平均值结果如图 3 所示
图 3
该基准测试在以下 4 种不同的设置下运行:
CPU(LLVM): 模型被编译为 LLVM IR 和 JITed, 完全在 CPU 上运行
OpenCL: 模型被编译成 OpenCL 还有一些 glue code 被编译到 LLVM, 负责设置和启动 OpenCL 内核然后我们在本地机器上运行
OpenGL: 与 OpenCL 相同, 但编译为 OpenGL
WebGL:glue code 被编译为 LLVM, 并使用 Emscripten 转换为 JavaScript 设备代码被编译为 WebGL 我们在 Firefox 上执行模型
从上面的结果可以看出, TVM OpenGL 后端与 OpenCL 具有相似的性能有趣的是, 浏览器的 WebGL 版本并不比桌面 OpenGL 慢很多考虑到主机代码是 JavaScript, 这很令人惊讶这可能是由于 Emscripten 生成了 asm.js, 可以在 Firefox 中显著优化
这是将深度学习模型自动编译到 Web 浏览器的第一步随着我们将优化带入 TVM 堆栈, 可以期待更多性能改进
TVM: 一个端到端的优化堆栈
可扩展框架, 如 TensorFlow,MXNet,Caffe 和 PyTorch 是目前深度学习领域最流行和实用的框架但是, 这些框架只针对范围较窄的服务器级 GPU 进行了优化, 如果要将工作负载部署到其他平台(例如手机, 嵌入式设备和专用加速器 FPGAASIC 等), 就需要大量费力的工作我们提出一个端到端的优化堆栈 TVM, 具备图形级和运算符级的优化, 以为不同硬件后端提供深度学习工作负载的性能可移植性我们讨论了 TVM 解决深度学习优化的挑战: 高级操作符融合跨线程的低级内存重用任意硬件基元的映射, 以及内存延迟隐藏实验结果表明, TVM 在多个硬件后端的性能可与现有支持低功耗 CPU 和服务器级 GPU 的最优库相媲美我们还通过针对基于 FPGA 的通用深度学习加速器的实验, 展示了 TVM 对新的硬件加速器后端的适应能力该编译器基础结构已经开源
图 4:CPUGPU 与 TPU 类的加速器需要不同的片上存储架构和计算基元在生成优化代码时我们必须考虑这个问题
我们提出 TVM, 一个端到端的优化编译器堆栈(如图 5 所示), 它能降低和微调深度学习的工作负载, 以适应多种硬件后端 TVM 的设计目的是分离算法描述 schedule 和硬件接口, 这个原则受到 Halide 的 compute/schedule 分离的想法的启发, 而且通过将 schedule 与目标硬件内部函数分开进行了扩展这一额外的分离可以支持新的专用加速器及其相应的新的内部函数
图 5:TVM 堆栈图目前的堆栈支持多种深度学习框架以及主流 CPUGPU 以及专用深度学习加速器
TVM 具有两个优化层: 一个是计算图优化层; 另一个是带有新的 schedule primitives 的张量优化层结合这两个优化层, TVM 可以从大多数深度学习框架中获取模型描述, 执行高级和低级优化, 并为后端生成特定硬件的优化代码, 例如 Raspberry Pi,GPU 和基于 FPGA 的专用加速器本研究的贡献如下:
我们构建了一个端到端的编译优化堆栈, 能够将高级框架 (包括 Caffe,MXNet,PyTorch,Caffe2,CNTK) 中指定的深度学习工作负载部署到多种硬件后端 (包括 CPU,GPU 和 基于 FPGA 的加速器, 以及今天发布的 OpenGL / WebGL 后端)
我们提出了在不同硬件后端中为深度学习工作负载提供性能可移植性的主要优化挑战, 并引入新的 schedule primitive, 以利用跨线程内存重用新硬件的内部函数和延迟隐藏
我们在基于 FPGA 的通用加速器上对 TVM 进行评估, 提供关于如何最优适应特定加速器的具体案例
我们的编译器可以生成可部署的代码, 其性能可与当前最优的库相媲美, 并且可适应新的专用加速器后端
图 6: 两层卷积神经网络的计算图示例图中的每个节点代表一次运算, 消耗一个或多个 tensor, 并产生一个或多个 tensor
更多细节, 请阅读论文: https://arxiv.org/pdf/1802.04799.pdf
XLA v.s. TVM: 深度学习中间表示之争
从结构上看, TVM 是一个完整的深度学习中间表示 (IR) 堆栈的基础层(base layer), 提供了一个可重用的工具链, 用于编译高级神经网络算法, 生成适合特定硬件平台的低级机器代码
借鉴构建编译器的方法, 团队构建了一个两级的中间层, 由 NNVM(用于任务调度和内存管理的高级 IR)和 TVM(优化计算内核的低级 IR) TVM 随附一套可重复使用的优化库, 能够随意调整, 适应从可穿戴设备到高端云计算服务器的各种硬件平台的需求
谷歌的 XLA 也是使用两层优化的结构, 但 XLA 只针对 TensorFlow 相比之下, TVM 则试图成为一个开放的接口
不仅如此, 陈天奇之前在回答知乎提问如何评价陈天奇的模块化深度学习系统 NNVM? 时表示, TVM 和已有的解决方案不同, 以 XLA 作为例子, TVM 走了和目前的 XLA 比更加激进的技术路线, TVM 可以用来使得实现 XLA 需要的功能更加容易
根据 TVM 博客, 现在已经做了如下更新:
在最近统计的深度学习开源框架排名中, TensorFlow 均位于第一, 毫无争议但未来, 中间表示 (IR) 将成为深度学习框架之间竞争的关键
来源: https://yq.aliyun.com/articles/534436