[論文レビュー] PyTorch: An Imperative Style, High-Performance Deep Learning Library

Adam Paszke, Sam Gross+19 more|arXiv (Cornell University)|Dec 3, 2019

Parallel Computing and Optimization Techniques被引用数 16,163

ひとこと要約

PyTorch は Python での eager, imperative 実行を GPU 加速と自動微分と組み合わせ、使いやすさを維持しつつ競争力のある性能を実現します。Pythonic で柔軟なモデル構築体験を強調し、優れた相互運用性と高性能コアを持ちます。

ABSTRACT

Deep learning frameworks have often focused on either usability or speed, but not both. PyTorch is a machine learning library that shows that these two goals are in fact compatible: it provides an imperative and Pythonic programming style that supports code as a model, makes debugging easy and is consistent with other popular scientific computing libraries, while remaining efficient and supporting hardware accelerators such as GPUs. In this paper, we detail the principles that drove the implementation of PyTorch and how they are reflected in its architecture. We emphasize that every aspect of PyTorch is a regular Python program under the full control of its user. We also explain how the careful and pragmatic implementation of the key components of its runtime enables them to work together to achieve compelling performance. We demonstrate the efficiency of individual subsystems, as well as the overall speed of PyTorch on several common benchmarks.

研究の動機と目的

• 命令型で動的な実行が静的グラフの性能と同等になりうることを、深層学習フレームワークで実証する。
• PyTorch における使いやすさと高速性の両立を実現する設計原則を説明する。
• Python エコシステムとの相互運用性とカスタム演算の拡張性を示す。
• 効率的な CPU/GPU 実行と autograd を可能にする実装上の選択を説明する。
• 他の主要なフレームワークと比較した実証的評価を提供する。

提案手法

• モデル・データ・学習ループが通常の Python プログラムとなる、Python 中心の eager 実行モデルを提示する。
• 制御フロー (Python/C++) とデータフロー (テンソルと演算) の分離、および CUDA ストリームによる非同期 GPU 実行を説明する。
• 割り当てオーバーヘッドと断片化を減らすカスタム CUDA メモリアロケータを説明する。
• 拡張された torch.multiprocessing によるデータ並列性と共有 CUDA テンソルを可能にするマルチプロセッシングのサポートを詳述する。
• Python の GC と統合された参照カウント型メモリ管理戦略を概説し、予測可能なメモリ解放を実現する。
• 非同期データフロー、メモリ管理、および他フレームワークとのベンチマークにわたる評価を提供する。

実験結果

リサーチクエスチョン

• RQ1 実行時に定義するスタイルを特徴とする命令型フレームワークは、静的グラフフレームワークと同等の性能を深層学習ワークロードで達成できるか？
• RQ2 PyTorch をどのような設計選択が、ユーザーフレンドリー（Pythonic）でありながら GPU 上で高性能を発揮させるのか？
• RQ3 効率を維持しつつ、PyTorch はより広い Python エコシステムやカスタム拡張とどのように統合されるか？
• RQ4 使いやすさを犠牲にせず性能を引き出す主なランタイム機構（メモリ管理、CUDA ストリーム、autograd）は何か？

主な発見

• PyTorch は一般的なベンチマークのセットで、最速の競合フレームワークに対して 17% の範囲内の性能を達成する。
• CUDA ストリームを介して Python の制御フローと GPU 作業を重ね合わせることで、高いデバイス利用率を実現する非同期 GPU 実行。
• カスタムキャッシュ済みアロケータと参照カウント型メモリ管理が CUDA 割り当てのオーバーヘッドを減らし、効率的なメモリ再利用を可能にする。
• NumPy および DLPack との双方向の相互運用性によりゼロコピーのデータ共有が可能になり、Python データサイエンスツールとの統合を促進する。
• データ並列学習を加速するため、テンソルデータを共有メモリに移動し、効率的なプロセス間通信を可能にするマルチプロセッシング拡張。
• 研究コミュニティ内での採用シグナルが強いことを示している（arXiv の言及）。