TensorFlowプロジェクトは、同ディープラーニングフレームワークのバージョン2.4.0リリースを発表した。CUDA 11とNVIDIAのAmpere GPUアーキテクチャのサポートに加えて、分散トレーニングのための新たなストラテジとプロファイリングツールを備えている。その他のAPIのアップデートには、Kerasでの混合精度演算(mixed-precision)とNumPyフロントエンドが含まれる。
TensorFlowチームは、リリースの概要をブログ記事で提供している。分散トレーニングでは、ParameterServerStrategyが試験レベルで新たに導入される一方で、MultiWorkerMirroredStrategyが試験レベルから安定機能に移行された。TensorFlow Profilerにも、MultiWorkerMirroredStrategyジョブのプロファイリングのサポートが含まれている。Keras APIでは、混合精度演算が試験レベルから安定機能に移行すると同時に、いくつかのベースクラスがリファクタリングされて、メモリ消費の改善やカスタムトレーニングロジックの簡素化が実現されている。
試験用に提供されるNumPy用のフロントエンドを使用すれば、TensorFlow上で動作するNumPyコードを記述して、TensorFlowの実行最適化を活用することが可能になる。2.4リリース用のバイナリはCUDA 11およびcuDNN 8で動作するようにビルドされ、NVIDIAの新しいAmpere GPUハードウェアをサポートする。
データ並列性(data parallelism)は、複数のワーカに対してトレーニングデータをスライスする、分散トレーニングのテクニックである。TensorFlowのtf.distributeモジュールは、複数のGPUを備える単一マシン、あるいは複数マシンからなるクラスタを対象とした、同期および非同期分散トレーニングシナリオ用に、いくつかの異なるデータ並列ストラテジをサポートする。
同期トレーニングでは、各ワーカが自身のデータのスライスから勾配(gradient)を独自に計算し、その結果を収集してモデルの更新を行う。新たなバッチを実行する前に、各ワーカは待機するか、あるいは同期する必要がある。今回のリリースで試験レベルから安定機能に移行したMultiWorkerMirroredStrategyは、複数のマシンにわたる同期分散トレーニングを実装しており、TensorFlow ProfilerがMultiWorkerMirroredStrategyトレーニングジョブをサポートするようにアップデートされている。
新リリースには、試験レベルのParameterServerStrategyも含まれており、パラメータサーバにホストされているモデルの非同期更新が可能になった。ワーカが相互に依存しないため、クラスタのフォールトトレランス性が向上する。
TensorFlow上にモデルを記述するためのハイレベルフレームワークであるKerasには、混合精度APIが含まれている。混合精度は、高速な16ビット演算をモデルの一部に使用し、その他の部分は数学的な安定性のために32ビット精度を残しておくことで、パフォーマンスを"3倍以上"に向上するものだ。新リリースではこの混合精度APIが、試験レベルから安定機能に移行されている。Keras Functional APIもリファクタされ、メモリ消費量が改善されると同時に、呼び出しロジックが簡略化されている。
新しい試験的モジュールであるtf.experimental.numpyは、広く使用されているNumPy科学計算パッケージのAPIのサブセットを備える。このモジュールは、内部的にTensorFlowのTensorオブジェクトをラップした、NumPyのNDArrayの実装を提供する。これにより、TensorFlowのランタイムアクセラレーションを活用したNumPy互換コードの記述が可能になる。このモジュールは試験レベルであるため、今後、非互換的な変更が実施される可能性がある。またドキュメントには、サポート予定のないNumPyの機能として、NumPy C API、Swigインテグレーション、Fortranストレージオーダなどが挙げられている。
今回のリリースのTensorFlowバイナリは、CUDA 11をターゲットとしてビルドされる。NVIDIA A100 Ampere GPUの新機能をサポートすることにより、これまでのV100 GPUに比較して、トレーニング時間の改善が実現されている。CUDA 11上のTensorFlow 2.4は、デフォルトでTensorFlow-32データ型のサポートが有効になっている。これにより、精度は低下するが、行列計算は高速になる。もうひとつの有力なディープラーニングフレームワークであるPyTorchは、10月にバージョン1.7をリリースしている。同じくCUDA 11とA100 GPUをサポートする他、Windowsオペレーティングシステムでの分散トレーニングオプションを備えている。
リリース発表への反応として、SlackのマシンラーニングエンジニアであるLeigh Johnson氏は、次のようにツイートしている。
Parameter Serer Trainingを早く試してみたくてたまりません。分散型の中断可能な(スポット)ワークロードのコーディネイトは、モデリングのコストや計算フットプリントに大きく影響する難しい問題なのです。