[論文レビュー] tf.data: A Machine Learning Data Processing Framework

Derek G. Murray, Jiřı Šimša+2 more|arXiv (Cornell University)|Jan 28, 2021

Advanced Data Storage Technologies参考文献 48被引用数 9

ひとこと要約

tf.data は、宣言的で合成可能なオペレータと自動的なパフォーマンスチューニングを備えた機械学習データ処理フレームワークであり、並列処理、パイipelニング、キャッシュ、静的最適化を適用することで、ResNet50 の ImageNet におけるトレーニング収束を最大 10.4× 速め、手動チューニングを必要とせずにエキスパートチューニングされたパイプラインと同等の性能を達成する。

ABSTRACT

Training machine learning models requires feeding input data for models to ingest. Input pipelines for machine learning jobs are often challenging to implement efficiently as they require reading large volumes of data, applying complex transformations, and transferring data to hardware accelerators while overlapping computation and communication to achieve optimal performance. We present tf.data, a framework for building and executing efficient input pipelines for machine learning jobs. The tf.data API provides operators which can be parameterized with user-defined computation, composed, and reused across different machine learning domains. These abstractions allow users to focus on the application logic of data processing, while tf.data's runtime ensures that pipelines run efficiently. We demonstrate that input pipeline performance is critical to the end-to-end training time of state-of-the-art machine learning models. tf.data delivers the high performance required, while avoiding the need for manual tuning of performance knobs. We show that tf.data features, such as parallelism, caching, static optimizations, and non-deterministic execution are essential for high performance. Finally, we characterize machine learning input pipelines for millions of jobs that ran in Google's fleet, showing that input data processing is highly diverse and consumes a significant fraction of job resources. Our analysis motivates future research directions, such as sharing computation across jobs and pushing data projection to the storage layer.

研究の動機と目的

• 機械学習ワークロードにおける非効率なデータ入力パイプラインが引き起こすパフォーマンスボトルネックを解消すること。
• 並列処理やプリフェッチバッファサイズといったパフォーマンスに重要なパrameter の手動チューニングの負担を軽減すること。
• 視覚、自然言語処理、強化学習を含む多様な機械学習ワークロードで、高パフォーマンスなデータ処理を可能にすること。
• スケールでの実世界のデータパイプライン動作を特定し、今後の最適化の方向性を示すこと。
• TensorFlow や PyTorch などの既存の機械学習フレームワークとシームレスに統合できる再利用可能で合成可能な API を提供すること。

提案手法

• 状態なしのデータセットと状態付きのイテレータを用いた宣言的 API を提供し、データ処理パイプラインを定義する。
• パイプラインをデータフロー・グラフとして表現し、グラフリライトによる静的最適化を適用して効率を向上させる。
• パフォーマンス予測を用いて、並列処理の度合いやプリフェッチバッファサイズなどの内部パrameter を自動的にチューニングする。
• ソフトウェアパイプライン、キャッシュ、オペレータの統合といった主要な最適化をサポートし、計算とデータロードを重ねて実行する。
• GPU や TPU などのハードウェアアクセラレータと統合し、CPU にデータ前処理をオフロードして I/O と計算を重ねて実行する。
• 処理時間とメモリ使用量を軽量なインストルメンテーションで監視し、動的スケーリングの意思決定を可能にする。

実験結果

リサーチクエスチョン

• RQ1 現代の機械学習ワークロードにおいて、入力パイプラインのパフォーマンスはエンドツーエンドのトレーニング時間にどのように影響するか？
• RQ2 並列処理やプリフェッチサイズといったパイプラインパrameter の自動チューニングは、エキスパートチューニングされた設定にどの程度近づけるか？
• RQ3 多様なワークロードとデータモダリティにおいて、実世界の機械学習データパイプラインの主なパフォーマンスボトルネックは何か？
• RQ4 冗長な計算やリソース使用を最小限に抑えることを前提に、スケールでのデータ処理ワークロードをどのように最適化できるか？
• RQ5 機械学習データパイプラインで高いスループットと低レイテンシを達成するために、最も効果的なシステムレベルの抽象化と最適化は何か？

主な発見

• 公開ベンチマークでは、入力パイプラインがトレーニング時間の最大 65% を占め、Google のファイアムでは合計計算時間の 30% がデータインジェストに費やされている。
• tf.data は並列処理とソフトウェアパイプラインのみで、ResNet50 のトレーニング収束時間を 10.4× に短縮している。
• キャッシュや静的グラフリライトといった追加の最適化により、トレーニング速度がさらに 2× 向上している。
• tf.data の自動チューニングは、手動設定を必要とせずにエキスパートチューニングされたパイプラインと同等のパフォーマンスを達成している。
• 数百万件の実際の ML ジョブでは、データ処理パターンの多様性が顕著であり、類似したパイプライン間で顕著な冗長計算が発生している。
• この分析は、ジョブ間での計算共有や、ストレージレイヤーにデータプロジェクションを押し上げて前処理のオーバーヘッドを低減するための今後の研究を促進する。