[论文解读] EagerPy: Writing Code That Works Natively with PyTorch, TensorFlow, JAX, and NumPy

Jonas Rauber, Matthias Bethge+1 more|arXiv (Cornell University)|Aug 10, 2020

Computational Physics and Python Applications参考文献 13被引用 4

一句话总结

EagerPy 是一个 Python 框架，可编写与框架无关的深度学习代码，原生支持 PyTorch、TensorFlow、JAX 和 NumPy，且无性能开销。它通过透明分发统一各框架的 API，支持方法链式调用，并提供全面的类型注解，使库只需实现一次即可在多个框架中使用，同时保持完整性能和类型安全。

ABSTRACT

EagerPy is a Python framework that lets you write code that automatically works natively with PyTorch, TensorFlow, JAX, and NumPy. Library developers no longer need to choose between supporting just one of these frameworks or reimplementing the library for each framework and dealing with code duplication. Users of such libraries can more easily switch frameworks without being locked in by a specific 3rd party library. Beyond multi-framework support, EagerPy also brings comprehensive type annotations and consistent support for method chaining to any framework. The latest documentation is available online at https://eagerpy.jonasrauber.de and the code can be found on GitHub at https://github.com/jonasrauber/eagerpy.

研究动机与目标

• 消除库开发者为 PyTorch、TensorFlow、JAX 和 NumPy 维护独立实现的需要。
• 让用户能够在不被锁定于单一库的情况下自由切换深度学习框架。
• 提供统一、高性能且类型安全的 API，同时保留底层框架的原生性能。
• 实现与框架特定代码的无缝集成，支持渐进式采用框架无关代码。
• 提升代码可维护性，减少在不同深度学习框架间重复的代码量。

提出的方法

• EagerPy 引入统一 API，将操作透明地映射到底层框架的原生张量操作，避免中间的 NumPy 转换。
• 通过直接将操作委派给原始框架的张量，保持原生性能，避免昂贵的 CPU-GPU 内存传输。
• 通过在张量对象上暴露所有操作为方法，支持方法链式调用，使代码如 x.square().sum().sqrt() 舒适且易读。
• 使用 MyPy 实现全面的类型注解，支持静态类型检查，可在函数参数和返回值处提前发现错误。
• EagerPy 通过为每种操作定义特定于框架的翻译层，支持多种后端（包括 PyTorch、TensorFlow 2、JAX 和 NumPy）。
• 框架采用测试驱动开发，通过端到端测试确保所有支持框架中行为一致和 API 正确。

实验结果

研究问题

• RQ1 是否可以构建一个单一、统一的 API，原生支持 PyTorch、TensorFlow、JAX 和 NumPy，且不造成性能下降？
• RQ2 如何在具有不同语法和语义规范的框架之间实现方法链式调用和一致的 API 设计？
• RQ3 全面的类型注解在多大程度上能提升深度学习库的代码安全性和可维护性？
• RQ4 框架无关代码能否在不进行完整重构的前提下，逐步集成到现有库中？
• RQ5 如何将 PyTorch、TensorFlow 和 JAX 的自动微分 API 抽象为一致、高层的接口？

主要发现

• EagerPy 通过直接委派给原生张量，实现了在 PyTorch、TensorFlow、JAX 和 NumPy 上原生运行的框架无关代码，且无性能开销。
• 该框架成功支持方法链式调用，使 x.square().sum().sqrt() 这类直观、流畅的代码在所有支持的后端中均可使用。
• 全面的类型注解已使用 MyPy 实现并验证，显著提升了代码安全性和开发人员信心。
• 该框架已被成功采用为 Foolbox Native 的基础，使单一代码库能够以原生性能和完整类型安全支持多个框架。
• EagerPy 的测试驱动开发方法确保了所有支持框架中行为一致和 API 正确，降低了回归风险。
• GUDHI 等库已采用 EagerPy，以无代码重复的方式在 PyTorch、TensorFlow 和 JAX 上添加自动微分支持。