AI绘画

绘制江湖儿女,绘制世间万物

MIT开源高性能主动微分框架Enzyme:速度提升4.5倍

麻省理工学院(MIT)提出的主动微分框架 Enzyme 在 NeurIPS 2020 大会上引起过不少人的兴趣,现在人们已经可以尝试运用这套对象了。

当前,PyTorch、TensorFlow 等呆板进修框架已经成为了人们开发的重要对象。计算反向传播、贝叶斯推理、不确定性量化和概率编程等算法的梯度时,我们需要把所有的代码以微分型写入框架内。这对于将呆板进修引入新领域带来了问题:在物理模拟、游戏引擎、气候模型中,原领域组件不是由呆板进修框架的特定领域谈话(DSL)编辑的。因此在将呆板进修引入科学计算时,重写需求成为了一个挑战。为了解决这一问题,现在的发展趋势包含构建新的 DSL,让重写过程变得简单,或者在编程时直接进行构建。这些方法可以让我们获得有效的梯度,但是仍然需要运用 DSL 或可微分的编程谈话进行重写。开发者们自然会希望运用呆板进修框架以外的代码重用已有对象,或在一种谈话里写好损失函数,然后为其用例提供更简单的抽象。尽管目前已经出现了一些针对不同谈话的反向主动微分框架(AD),但我们很难在 ML 框架外的代码上主动运用它们,因为其仍需要重写,且对于跨谈话 AD 和库的支持有限。为了方便开发者,来自 MIT 的研究者开源了 Enzyme,一种用于 LLVM 编译器框架的高性能主动微分(AD)编译器插件。该插件能够分解以 LLVM 中间表示(IR)表示的固态可分析程序的梯度。Enzyme 能够分解任何以面向 LLVM IR 编译器为谈话编辑的程序的梯度,包括 C、C ++、Fortran、Julia、Rust、Swift、MLIR 等,从而提供这些谈话的本机 AD 功能。据作者介绍,与传统的源到源和 operator-overloading 对象不同,Enzyme 在优化的 IR 上执行 AD。在包括微软 ADBench 在内的以呆板进修为重点的基准套件上,经过优化的 IR 上的 AD 的几何平均速度比未经过优化的 IR 上 AD 的几何平均速度提高了 4.5 倍,这使得 Enzyme 达到了最高的性能。此外,Enzyme 方便运用,在 PyTorch 和 TensorFlow 上都有程序包,可让开发者便捷访问具有最新性能的外来代码梯度,从而使外来代码可直接合并到现有的呆板进修工作流程中。

项目网页:https://enzyme.mit.edu/

项目地址:https://github.com/wsmoses/Enzyme

作者表示,Enzyme 可提供这些对象和能力:

Enzyme,一种用于 LLVM 的编译器插件,可以分解可固态微分的 LLVM IR 的快速梯度。包括 C、C ++、Fortran、Rust、Swift 等编译器前端生成的中间表示(IR)。

PyTorch-Enzyme/TensorFlow-Enzyme,一个外部功能接口,允许呆板进修研究者运用 PyTorch 和 TensorFlow 运用以 LLVM 编译谈话编辑的外部代码。

Enzyme.jl,一个 Julia 包,通过动态高级谈话编辑的代码,仅运用低层信息获得梯度。

通过链接时优化(LTO)支持了多来源 AD 和固态库支持。

研究表明优化后运行 AD,在标准呆板进修基准测试上可获得显著性能提升,并达到 SOTA 水平。

上图为 relu(pow(x,3)) 的梯度分解示例。左侧为 LLVM IR 上的原始计算。左侧注释中展示了将添加到前向传递中的活动变量的影子分配。右侧则是 Enzyme 将生成的反向传递。完整的分解梯度函数将结合运用这些函数(添加影子分配),将 if.end 中的返回替换为 reverse_if.end 的分支。Enzyme 项目概述Enzyme 项目是一个用于可固态分析 LLVM IR 的反向模式主动微分(AD)对象。它允许开发者可以主动创建基于源代码的梯度,而无需更多额外工作。double foo(double);double grad_foo(double x) { return __enzyme_autodiff(foo, x);}通过优化微分后的代码,Enzyme 可以比现有的优化对象提供更快的微分速度:

组件Enzyme 对象主要由四部分组成:

可选的预处理阶段,该阶段执行对 AD 有用的较小转换。

一种新的过程间类型分析,可推断出内存位置的基础类型。

活动分析,确定哪些指令或值会影响导数计算(在现有 AD 系统中很常见)。

优化遍历可创建任何必需的派生函数,用生成的函数替换对__enzyme_autodiff 的调用。

更多详细介绍,可查看 MIT 研究者们提交的 NeurIPS 2020 论文:

论文地址:https://arxiv.org/pdf/2010.01709.pdf

Powered by Froala Editor

留下评论

您的电子邮箱地址不会被公开。 必填项已用*标注