FAIR 开源 Tensor Comprehensions，让机器学习与数学运算高性能衔接

AI研习社

发布于 2018-03-16 10:59:36

8710

发布于 2018-03-16 10:59:36

AI 研习社消息，Facebook AI 研究院于近日开源了 C++ 库及数学语言 Tensor Comprehensions，它能有效填补研究人员于数学运算领域的沟通鸿沟，并基于各种硬件后端上大规模运行工程模型。它的最大亮点在于，它采用了 Just-In-Time 的编译自动生成机器学习社区所需的高性能代码。

也就是说，通过 Tensor Comprehensions，研究人员能够以数学符号的方式进行编写，系统能够根据需求进行编译调整，并输出专业的代码。

数量级增长

传统意义上，如果要从零创造一个具有高性能表现的机器学习层需要两个大步骤，这可能需要耗费工程师数天甚至数周的努力。

1. 在 NumPy 层级，研究人员需要写一个全新层，并在以 PyTorch 为代表的深度学习库中链接已有运算，然后进行小规模测试。如果要运行大规模实验，需要通过数量级加速检测其代码实现。
2. 工程师随后采用这个层，并为 GPU 和 CPU 撰写有效代码（但这需要满足多个条件）： 1）工程师需要是一位在高性能计算颇有了解的专家，但这一人才目前非常稀缺； 2）文本、策略、写代码，debug，工程师需要样样精通； 3）将代码与实际任务相关的后端相连接，如冗长的参数检查和添加样板集成代码

这也直接导致近年来深度学习社区一直依赖以 CuBLAS, MKL, 和 CuDNN 为代表的高性能库而构建运行于 GPU 和 CPU 的代码。对于研究员而言，如何寻找一条新的实现思路成为了巨大挑战。

而在 Tensor Comprehensions 中，研究员们得以将这一流程从几天甚至几周减少到数分钟。这一开源包含了：

用简单语法表达一系列机器学习概念的数学符号
基于 Halide IR 数学符号的 C ++前端
基于整数集库（ISL）的 Just-in-Time 编译器，
一个基于进化搜索的多线程、多 GPU 自动调节器

早期工作

Halide 是一种最近在高性能图像处理领域颇受欢迎的语言，它采用类似的高级函数语法来描述一个图像处理的 pipeline，随后在单独代码块中调度到硬件上，并且详细到如何平铺、矢量化、并行化和融合。对于具有专业知识的人而言，这是一种非常高效的语言；但对于机器学习从业者来说，这一难度并不小。Halide 的自动调度在研究上非常活跃，但对于 GPU 上运行的机器学习代码，目前还没有很好的解决方案。

Tensor Comprehension 将 Halide 编译器作为调用库。FAIR 研究员构建了 Halide 的中间表征（IR）和分析工具，并与多面编译进行技术配对，因此，用户可以在无需理解运行原理的条件下使用相似的高级语法编写层。此外，FAIR 研究员也找到了简化语言的方法，不需要为缩减运算制定循环边界。

细节实现

使用 Halide 和多面编译技术，Tensor Comprehensions 能通过委托内存管理和同步功能以自动合成 CUDA 内核。这一编译能够针对特定尺寸对一般操作符进行融合，对快速本地内存、快速缩减和 Just-in-Time 专业化都能优化。因不尝试涉足内存管理，因此这一流程能够轻松集成到机器学习的任意框架，以及任何允许调用 C++ 函数的语言中。

与经典编译器技术和库方法相反的是，多面编译能让张量理解为每个新网络按需调度单个张量元素的计算。

在 CUDA 层面，它结合了仿射循环转换、融合/分裂和自动并行处理，同时确保数据在存储器层次结构中正确移动。

图中数字表示最初计算张量元素的顺序，箭头表示它们之间的依赖关系。在这个例子中，数字旋转对应于循环交换，能够实现深度操作器的融合。

研究员提供了一个集成的多线程、多 GPU 自动调节库，以推进搜索过程，它使用 Evolutionary Search 来生成和评估数千种实现方案，并选择性能最佳的方案。只需调用 Tensor Comprehension 上的 tune 功能，用户便能目睹性能的改善。最好的策略是通过 protobuf 序列化，在当下或离线情况下可重用。

在性能方面，Tensor Comprehensions 可能仍然有很大提升空间，但在条件满足的情况下，它已经可以匹配或打破当前机器学习框架与手动调整库集成的性能，并通过将代码生成策略适应特定问题大小的能力来实现。下图展现了 Tensor Comprehensions 与 Caffe2 和 ATen 的对比。