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

雷锋网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 的对比。

更多信息请参考 arXiv 论文（链接见底部）。

未来，Tensor Comprehensions 将补充硬件制造商（如 NVIDIA 和 Intel）编写的快速库，并将与 CUDNN，MKL 或 NNPack 等库一起使用。

下一步

雷锋网AI科技评论了解到，该版本的开源库能让研究人员和程序员使用与他们在论文中使用的数学相似标记来编写层，并简明地传达程序意图。他们也能够在几分钟内将这种表示方法进行编译。未来随着工具链的不断增长，FAIR 认为它的预计可用性和性能将会提升，并受益整个社区。

晚些时间，FAIR 将会发布 PyTorch 集成的 Tensor Comprehensions。

GitHub 链接：https://facebookresearch.github.io/TensorComprehensions/

arXiv 论文：https://arxiv.org/abs/1802.04730

Slack 页面：https://tensorcomprehensions.herokuapp.com/

Email: tensorcomp@fb.com

雷锋网AI科技评论编译