[論文レビュー] 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 は、パフォーマンスのオーバーヘッドなしに PyTorch、TensorFlow、JAX、NumPy とネイティブに連携するフレームワークに依存しないディープラーニングコードを書くための Python フレームワークです。透明なディスpatch を通じて 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 用に別々の実装を維持する必要をなくすこと。
• ユーザーが1つのライブラリに縛られず、ディープラーニングフレームワークを自由に切り替えられること。
• 下位のフレームワークのネイティブパフォーマンスを保ちながら、統一的でパフォーマンスが高く、型安全な API を提供すること。
• 段階的な導入を可能にするために、フレームワークに依存しないコードとフレームワーク固有のコードをシームレスに統合できること。
• ディープラーニングフレームワーク間でのコードの保守性を向上させ、重複を削減すること。

提案手法

• EagerPy は、中間の NumPy 変換を回避するように、操作を下位のフレームワークのネイティブテンソル操作に透明にマップする統一された API を導入しています。
• CPU と GPU 間のメモリ転送を回避することで、元のフレームワークのテンソルに直接操作をデリゲートすることで、ネイティブパフォーマンスを維持しています。
• メソッドチェイニングを実現するために、すべての操作をテンソルオブジェクトのメソッドとして公開しており、x.square().sum().sqrt(). のような直感的で読みやすいコードが可能になっています。
• MyPy を用いた包括的な型アノテーションにより、関数のパラメータと戻り値の両方で静的型チェックが可能になり、バグの早期検出が可能になっています。
• PyTorch、TensorFlow 2、JAX、NumPy の各バックエンドをサポートするために、各操作ごとにフレームワーク固有の翻訳レイヤーを定義しています。
• エンドツーエンドのテストを用いたテスト駆動開発により、すべてのサポート対象フレームワークで一貫した振るまいと API の正しさを保証しています。

実験結果

リサーチクエスチョン

• RQ1 パフォーマンスの低下なしに、PyTorch、TensorFlow、JAX、NumPy をネイティブにサポートする単一の統一 API を構築できるか？
• 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 で自動微分のサポートが追加されました。