TensorFlow 2.9が利用可能になりました。CPUパフォーマンスの最適化を実現するためにoneDNNが改良され、WSL2がすぐに利用可能になりました。

昨日、TensorFlow は正式に発表しました: TensorFlow 2.9 がリリースされました!前回からの距離バージョン 2.8 アップデートはわずか3か月前に行われました。

新バージョンのハイライトは次のとおりです。

oneDNN のパフォーマンスの向上。
データ並列処理からモデル並列処理へのシームレスな移行を実現する新しい API である DTensor のリリース。
Eigen、tf.function の統合、Windows 上の WSL2 の新しいサポートを含むコアライブラリの改善。
tf.function の再トレースおよび Keras オプティマイザー用の新しい実験的な API もリリースされました。

TensorFlow 2.9 の具体的な改善点を見てみましょう。

TensorFlow 2.9 の新機能

CPU パフォーマンスの向上: oneDNN

TensorFlow は Intel と連携して、oneDNN パフォーマンスライブラリを TensorFlow に統合し、Intel CPU で最高のパフォーマンスを実現しました。 TensorFlow 2.5 以降、TensorFlow は oneDNN を実験的にサポートしており、最大 4 倍のパフォーマンス向上を実現できます。 TensorFlow 2.9 では、Linux x86 パッケージおよび Intel Cascade Lake 以降の CPU に搭載されている AVX512_VNNI、AVX512_BF16、AMX などのニューラルネットワークハードウェア機能を備えた CPU で、oneDNN 最適化がデフォルトで有効になっています。

Intel は次のように語っています。「Google AI と協力して、oneDNN ライブラリを TensorFlow 2.9 のデフォルトのバックエンド CPU 最適化にできることを嬉しく思います。」 TensorFlow ユーザーは、コードを変更することなく、オープンソフトウェアの最適化によって実現される AI パフォーマンスの高速化を実感できるようになります。

oneDNN 最適化を有効にして TensorFlow を実行しているユーザーは、最適化をオフにしたときとは若干異なる数値結果が見られる場合があります。これは、浮動小数点の丸め方法と順序が異なると、若干の不正確さが生じる可能性があるためです。

上記のことで問題が発生する場合は、TensorFlow プログラムを実行する前に TF_ENABLE_ONEDNN_OPTS=0 を設定して最適化をオフにしてください。再度有効にしたい場合は、TensorFlow プログラムを実行する前に TF_ENABLE_ONEDNN_OPTS=1 を設定します。最適化が有効になっていることを確認するには、プログラムログで「oneDNN custom operations are on」で始まるメッセージを探します。

テンソル

DTensor は、分散モデル処理用の新しい TensorFlow API であり、モデルをデータ並列処理から単一プログラム複数データ (SPMD) ベースのモデル並列処理にシームレスに移行できます。

DTensor のコア設計原則は次のとおりです。

デバイスに依存しない API: これにより、デバイスタイプ間で分割されたモデルを含め、CPU、GPU、または TPU で同じモデルコードを使用できるようになります。
マルチクライアント実行: コーディネーターを削除し、各タスクがローカルに接続されたデバイスを駆動できるようにすることで、起動時間に影響を与えずにモデルを拡張できます。
グローバルな視点とレプリカごとの視点: 従来の TensorFlow では、分散モデルコードはレプリカを中心に記述されていましたが、DTensor では、モデルコードはグローバルな視点から記述され、各レプリカコードは DTensor ランタイムによって生成および実行されます。

TensorFlow では、DTensor に関する入門チュートリアルをいくつか公式に開発しています。参考資料は次のとおりです。

DTensor の概念: https://www.tensorflow.org/guide/dtensor_overview
DTensor を使用した分散 ML: https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial
DTensor と Keras を一緒に使用する: https://www.tensorflow.org/tutorials/distribute/dtensor_keras_tutorial

tf.function の TraceType

新しいバージョンでは、tf.function の再トレース方法が改善され、よりシンプルで予測可能、かつ構成可能になりました。

tf.function へのすべての引数には、tf.types.experimental.TraceType が割り当てられます。カスタムユーザークラスは、トレースプロトコル (tf.types.experimental.SupportsTracingProtocol) を使用して TraceType を宣言できます。 TraceType システムを使用すると、トレースバックルールを簡単に理解できます。

WSL2のサポート

WSL2 を使用すると、開発者は従来の仮想マシンやデュアルブート設定を必要とせずに、Windows 上で直接 Linux 環境を実行できます。 TensorFlow は、GPU アクセラレーションを含め、WSL2 をすぐにサポートするようになりました。

決定論的行動

API tf.config.experimental.enable_op_determinism は TensorFlow の op を決定論的にします。決定論とは、ユーザーが同じ入力でオペレーションを複数回実行した場合、オペレーションが毎回まったく同じ出力を返すことを意味し、これはモデルのデバッグに役立ちます。

一般に、多くの操作は、非決定的な順序で浮動小数点数を加算できるスレッドを使用するため、非決定的です。

TensorFlow 2.8 では、オペレーションを決定論的にする API が導入され、TensorFlow 2.9 では、いくつかのケースで決定論的なパフォーマンスが向上しています。 TensorFlow モデルを確定的に実行したい場合は、プログラムの先頭に次のコードを追加するだけです。

 < コード> tf.keras.utils.set_random_seed ( 1 ) < / コード>
 < code > tf . config . experimental . enable_op_determinism ( ) </code>

最初の行は、決定論に必要な Python、NumPy、TensorFlow のランダムシードを設定します。 2 行目は、すべての TensorFlow 操作を決定論的にします。決定論は通常、パフォーマンスの低下を伴いますので、op 決定論を有効にするとモデルの実行速度が低下する可能性があることに注意してください。

Keras によるトレーニングの最適化

TensorFlow 2.9 では、Keras Optimizer API の新しい実験バージョンである tf.keras.optimizers.experimental が追加されました。 API は、ユーザーがより簡単にカスタマイズおよび拡張できる、より統合され拡張された組み込みオプティマイザーのカタログを提供します。

将来のリリースでは、tf.keras.optimizers.experimental.Optimizer (およびサブクラス) が tf.keras.optimizers.Optimizer (およびサブクラス) に置き換えられます。つまり、古い Keras オプティマイザーを使用するワークフローは、新しいオプティマイザーに自動的に切り替わります。

新しいオプティマイザークラスのハイライトは次のとおりです。