Facebook 发布张量理解库,自动编译高性能机器学习核心
   来源:量子位     2018年02月16日 03:43

Facebook AI Research 今天发布了张量理解(Tensor Comprehension),这是一个 C ++ 库,也是一种数学语言,它能够自动、按需地及时编译出机器学习所需的高性能代码。

简单来说,就是将(部分)人类能读懂的高级编程语言代码,转换成 CUDA 等专为机器优化的代码。

Facebook 在博客中说,Tensor Comprehension 有助于弥合研究人员与工程师之间的沟通鸿沟。研究人员往往专注于数学运算,而工程师专注于在多种硬件后端上运行大规模模型的实际需求。

以前,构建高性能机器学习层通常需要几天到几周的工作,分为两个环节:

一、研究人员在 numpy 级别的抽象中写一个新层,与 PyTorch 等深度学习库链接起来,进行小规模测试。要把这些用来验证想法的代码用到大规模实验上,需要加快一个数量级。

二、工程师根据这个层为 GPU 和 CPU 写高效代码。

这个过程中需要的工程师,得是高性能计算专家,这类人才数量非常有限。他们需要了解情境、制定策略、编写并调试代码,同时,这个将代码移到后端的工作还涉及一些杂活,比如冗长的参数检查、添加样板集成代码等。

因此,过去几年里,深度学习界在生成 GPU 和 CPU 尚运行的高性能代码时,越来越依赖 CuBLAS、MLK、CuDNN 等代码库。

当想要试验的新想法偏离了这些代码库提供的基元时,会涉及到大量的工程工作,研究人员会被吓得止步不前。

Facebook 认为,简化从想法到高性能代码的过程,把它从几天、几周缩短到几分钟有着重大的实践价值。他们希望研究人员从数学的角度描述自己的想法之后,通过 Tensor Comprehension 的自动编译和调整,就能生成性能良好的专用代码。

这一版本的 Tensor Comprehension 包含:

能用简单语法表达一系列机器学习想法的数学符号;

基于 Halide IR 的,表达这些数学符号的 C++ 前端;

一个基于整数集库(ISL)的多面即时(JIT)编译器;

一个基于进化搜索的多线程、多 GPU 自动调节器。

Tensor Comprehension 将高性能图像处理领域的流行语言 Halide 的编译器作为一个库,建立在 Halide 的中间表示(intermediate representation,简称 IR)基础上,和多面编译(polyhedral compilation)技术结合起来。让用户可以用类似的高级语法编写层,但不用说明它如何运行。

通过 Halide IR 和 polyhedral compilation 的结合,Tensor Comprehension 通过委托内存管理和同步功能自动合成 CUDA 核。这种翻译针对特定的运算符融合、快速本地内存、快速缩减和特定尺寸的 JIT 专门化进行优化。

Facebook 在博客中表示,因为这一工作流不尝试拥有或优化内存管理,因此可以轻松高效地集成到任何 ML 框架和能调用 C++ 函数的语言中。

不同于传统的编译器技术和库方法,多面编译让 Tensor Comprehension 能为每个新网络按需调度单个张量元素的计算。

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

图中的数字表示张量元素最初计算的顺序,箭头表示它们之间的依赖关系,数字旋转对应着循环交换,让深层运算符融合成为可能。

Facebook 还为推动搜索过程提供了一个集成的多线程、多 GPU 自动调整库,用进化搜索来生成和评估数千种实现方案,并选择性能最佳的方案。调用 Tensor Comprehension 上的调节(tune)功能就能实时看到性能的提升。

在性能方面,Tensor Comprehension 在部分案例中可以匹敌甚至超越集成了手动调节代码库的机器学习框架。这主要依靠让代码生成策略适应特定问题大小的能力。

下图可以看出,Tensor Comprehensions 自动生成核的性能,比 Caffe2 和 ATen 用供应商提供的 CuDNN 等库生成的核有所提升。

对于英伟达、英特尔等硬件厂商所提供的快速库来说,Tensor Comprehension 是一种补充,也可以和 CuDNN、MKL、NNPack 等一起使用。

稍后,Facebook 还将发布 Tensor Comprehension 的 PyTorch 集成。

Tensor Comprehensions 是 Facebook,法国国家信息与自动化研究所(Inria),苏黎世联邦理工学院和麻省理工学院的合作项目。

论文:

Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, Albert Cohen

https://arxiv.org/abs/1802.04730

代码:

https://github.com/facebookresearch/TensorComprehensions

文档:

https://facebookresearch.github.io/TensorComprehensions/

博客:

https://research.fb.com/announcing-tensor-comprehensions/

— 完 —

加入社群

量子位 AI 社群 13 群开始招募啦,欢迎对 AI 感兴趣的同学,加小助手微信 qbitbot5 入群;

此外,量子位专业细分群 ( 自动驾驶、CV、NLP、机器学习等 ) 正在招募,面向正在从事相关领域的工程师及研究人员。

进群请加小助手微信号 qbitbot5,并务必备注相应群的关键词 ~ 通过审核后我们将邀请进群。(专业群审核较严,敬请谅解)

诚挚招聘

量子位正在招募编辑 / 记者,工作地点在北京中关村。期待有才气、有热情的同学加入我们!相关细节,请在量子位公众号 ( QbitAI ) 对话界面,回复 " 招聘 " 两个字。

量子位 QbitAI · 头条号签约作者

վ ' ᴗ ' ի 追踪 AI 技术和产品新动态

超越 代码 张量