Cloud TPU VM でモデルをプロファイリングする

プロファイルにより、Cloud TPU でのモデルのトレーニング パフォーマンスを最適化できます。TensorBoardCloud TPU TensorBoard プラグインを使用して、モデルをプロファイリングします。

サポートされているフレームワークのいずれかと TensorBoard を使用する方法については、次のドキュメントをご覧ください。

トレーニング スクリプトのプロファイリングの前提条件

TPU プロファイリング ツールを使用する前に、次のことを行う必要があります。

  1. モデル トレーニング セッションを開始する

    1. v4-8 TPU を設定してモデルをトレーニングします。このドキュメントで説明するプロファイリング手順では ResNet モデルを使用しますが、v4 TPU でトレーニングする場合は別のモデルを使用できます。
    2. TPU VM で、トレーニング スクリプトにプロファイラ サーバーを起動する行を追加します。

      ResNET トレーニングの場合、トレーニング スクリプトは /usr/share/tpu/tensorflow/resnet50_keras/resnet50.py にあります。

      ハイライト表示された行を resnet50.py に挿入します。 ファイルの先頭に次のインポートを追加します。

      import tensorflow.compat.v2 as tf2
      

      スクリプトがトレーニング ループを開始する直前に、ハイライト表示された行を追加します。

      if name == 'main':
       tf.logging.set_verbosity(tf.logging.INFO)
       tf2.profiler.experimental.server.start(6000)
       app.run(main)
      

      TensorFlow プロファイラ サーバーは、スクリプトの実行時に TPU VM 上で起動します。

    3. モデルのトレーニングを開始します。

      トレーニング スクリプトを実行し、モデルがアクティブにトレーニングされていることを示す出力が表示されるまで待ちます。出力はコードやモデルによって異なります。Epoch 1/100 のような出力を探します。 または、Google Cloud コンソール の Cloud TPU ページに移動して TPU を選択し、CPU 使用率のグラフを表示することもできます。CPU 使用率のグラフには TPU 使用率は表示されませんが、TPU によってモデルがトレーニングされていることを示す良いサインです。

モデルのトレーニングのプロファイリングを開始する

モデルをトレーニングするときに、別のターミナル ウィンドウまたは Cloud Shell を開きます。モデルのトレーニングのプロファイリングを開始するには、次の手順を行います。

  1. 新しいウィンドウまたはシェルで、ポート転送を使用して TPU VM に接続します。

    gcloud compute tpus tpu-vm ssh your-vm --zone=us-central2-b --ssh-flag="-4 -L 9001:localhost:9001"

    ポート転送では、ローカル ブラウザが TPU VM 上で実行されている TensorBoard サーバーと通信できます。

  2. TensorFlow の要件 { id="install-tensorboard"} をインストールします。

    TPU VM には、デフォルトで TensorBoard がインストールされています。TensorFlow を手動でインストールすることもできます。どちらの場合も、追加の依存関係が必要になる場合があります。次のコマンドを実行して、TPU VM にこれらの依存関係をインストールします。

    pip3 install -r /usr/share/tpu/models/official/requirements.txt
  3. Cloud TPU TensorBoard プラグイン {: id="install-plugin"} をインストールします。

    TPU VM から、次のコマンドを実行します。

     pip3 install --upgrade "cloud-tpu-profiler>=2.3.0"
     pip3 install tensorflow
     pip3 install tensorboard_plugin_profile
  4. TensorBoard サーバーを起動する

    TensorBoard を実行し、TPU VM にログ ディレクトリ(logdir)を作成します。ここで、TensorBoard がプロファイリング データを書き込むことができます。--logdir フラグを使用してログ ディレクトリを指定します。次に例を示します。

    mkdir log-directory
    TPU_LOAD_LIBRARY=0 tensorboard --logdir log-directory --port 9001

TensorBoard でウェブサーバーが起動され、その URL が表示されます。

Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.3.0 at http://localhost:9001 (Press CTRL+C to quit)

ウェブブラウザを開き、TensorBoard の出力に表示されている URL に移動します。TensorBoard ページの右上にあるプルダウン メニューから [プロファイル] を選択します。使用可能なプロファイリング ツールのリストが、左側のサイドバーの [ツール] プルダウン メニューに表示されます。

イメージ

TPU VM でプロファイルをキャプチャする

  1. [CAPTURE PROFILE] ボタンを選択します。
  2. [IP アドレス] ラジオボタンをオンにします。
  3. Profile Service URL フィールドに「HOSTNAME:6000」と入力します。
  4. CAPTUREボタンを選択します。

イメージ

TensorBoard でプロファイル データを表示する

プロファイルをキャプチャすると、TensorBoard に overview_page が表示されます。使用可能なプロファイリング ツールのリストが左側のペインに表示されます。

イメージ

プロフィール

[Profile] タブは、一部のモデルデータがキャプチャされると表示されます。TensorBoard ページでの [更新] ボタンのクリックが必要になる場合があります。データが利用可能になり、[Profile] タブをクリックすると、パフォーマンス分析に役立つ次のようなツールが表示されます。次のいずれかのツールを使用して、モデルをプロファイリングできます。

プロファイルの概要ページ

[プロファイル] ページにある概要ページ(overview_page)には、キャプチャ中のモデルがどのように実行されたかが表示されます。このページには、すべての TPU を集計した概要と入力パイプライン全体の分析結果が表示されます。ホスト プルダウンから個別の TPU を選択することもできます。

このページでは、次のパネルにデータが表示されます。

イメージ

  • パフォーマンス サマリー

    • FLOPS 使用率 - TPU マトリックス ユニットの使用率
  • TPU での上位 10 個の TensorFlow 演算では、多くの時間が費やされた TensorFlow 演算が表示されます。

    各行は、演算に費やされた自己時間(すべての演算にかかった時間に占める割合)、累積時間、カテゴリ、名前、FLOPS 率を示します。

  • 実行環境

    • 使用されたホストの数
    • 使用された TPU のタイプ
    • TPU コアの数

入力パイプライン分析ツール

入力パイプライン分析ツールを使用すると、パフォーマンスの結果をより詳しく分析できます。 このツールは、プログラムで負荷の高い入力処理をすぐに通知し、デバイス側とホスト側で分析を行い、ボトルネックとなっているパイプラインのステージをデバッグできます。

パイプラインのパフォーマンスの最適化について詳しくは、入力パイプラインのパフォーマンスに関するガイダンスをご覧ください。

入力パイプライン

TensorFlow プログラムがファイルからデータを読み取ると、読み取りプロセスは連続した複数のデータ処理ステージに分割されます。1 つのステージの出力は、次のステージへの入力です。この読み込み方式を入力パイプラインといいます。

ファイルからレコードを読み取るための一般的なパイプラインには、次のステージがあります。

  1. ファイルの読み取り
  2. ファイルの前処理(オプション)
  3. ホストマシンからデバイスへのファイル転送

入力パイプラインが非効率な場合、アプリケーションの速度が大幅に低下する可能性があります。入力パイプラインに多くの時間が費やされている場合、このアプリケーションは入力バウンドとみなされます。入力パイプライン分析ツールを使用すると、非効率な入力パイプラインを特定できます。

入力パイプライン ダッシュボード

入力パイプライン分析ツールを開くには、[プロファイル] を選択し、[ツール] プルダウンから [input_pipeline_analyzer] を選択します。

ダッシュボードには、デバイス側とホスト側の分析の詳細が表示されます。

デバイス側の分析 - デバイスのステップ時間の詳細が表示されます。

  • デバイスのステップ時間の統計情報。
  • 入力データを待機しているデバイス ステップ時間の割合

ホスト側の分析

このセクションでは、ホスト側の分析の詳細をいくつかのカテゴリに分けて説明します。

  • デバイスに転送するデータのエンキュー。デバイスへの転送前にデータがインフィード キューに追加される際に費やされた時間。
  • データの前処理: 画像の圧縮など、前処理操作に費やされた時間。
  • ファイルからの事前データの読み取り: キャッシュ、プリフェッチ、インターリーブなど、ファイルの読み取りに費やされた時間。
  • ファイルからのオンデマンドのデータ読み取り: キャッシュ、プリフェッチ、インターリーブなしで、ファイルからデータを読み取る際に費やされた時間。
  • その他のデータの読み取りまたは処理tf.data を使用しない、その他の入力関連のオペレーションに費やされた時間。

イメージ

個々の入力演算の統計とそのカテゴリの内訳を実行時間別に表示するには、Show Input Op statistics セクションを展開します。

ソースデータ テーブルは次のようになります。

イメージ

各テーブル エントリには、次の情報が含まれます。

  1. 入力演算: 入力演算の TensorFlow 演算名が表示されます。
  2. カウント: プロファイリング期間中に実行された演算のインスタンスの合計数が表示されます。
  3. 合計時間(ミリ秒): 各オペレーション インスタンスに費やされた時間の累積合計が表示されます。
  4. 合計時間(%): 演算に費やされた合計時間が、入力処理に費やされた合計時間との割合で表示されます。
  5. 合計自己時間(ミリ秒): 関数のすべてのインスタンスの累積時間を表示します。この自己時間は、関数本体内で費やされた時間を測定したもので、関数本体から呼び出される関数で費やされた時間は含まれません。たとえば、Iterator::PaddedBatch::Filter::ForeverRepeat::MapIterator::PaddedBatch::Filter によって呼び出されるため、その合計自己時間は、後者の合計自己時間から除外されます。
  6. 合計自己時間(%): 合計自己時間が、入力処理に費やされた合計時間との割合で表示されます。
  7. カテゴリ: 入力演算の処理カテゴリを示します。

演算プロファイル

演算プロファイルは、プロファイリング期間中に実行された XLA 演算のパフォーマンス統計を表示する Cloud TPU ツールです。オペレーション プロファイルには次の情報が表示されます。

  • アプリケーションによる Cloud TPU の使用状況。演算に費やされた時間の割合がカテゴリ別に表示されます。また、TPU FLOPS の使用率も表示されます。
  • 最も時間のかかる演算。これらの演算は、最適化の潜在的なターゲットとなります。
  • 個々の演算の詳細。演算で使用される形状、パディング、式などが表示されます。

演算プロファイルを使用すると、最適化のターゲットを見つけることができます。たとえば、オペレーション プロファイルを使用して、実行に最も時間を費やした XLA 演算とその演算による TPU FLOPS の使用量を特定できます。

演算プロファイルの使用

演算プロファイル ツールには、XLA 演算のパフォーマンス統計が含まれています。画面上部の [プロファイル] タブをクリックし、[ツール] プルダウンから [op_profile] を選択すると、TensorBoard で演算プロファイル データを表示できます。次のような画面が表示されます。

イメージ

  1. 概要セクション: Cloud TPU の使用率を表示され、最適化のための推奨事項が表示されます。
  2. コントロール パネル: テーブルに表示されるオペレーションの数、表示されるオペレーション、並べ替え方法を設定するためのコントロールが含まれます。
  3. 演算テーブル: XLA 演算に関連付けられている TensorFlow 演算の上位カテゴリが一覧表示されます。これらの演算は、Cloud TPU 使用率の順に並んでいます。
  4. 演算詳細カード: テーブル内の演算にカーソルを合わせると、演算の詳細が表示されます。これらの詳細には、FLOPS 使用率、オペレーションが使用されている式、オペレーション レイアウト(fit)などがあります。

XLA 演算テーブル

演算テーブルには、Cloud TPU 使用率の高い順に XLA 演算カテゴリが表示されます。テーブルには、実行時間の割合、演算カテゴリ名、関連する TensorFlow 演算名、カテゴリの FLOPS 使用率が表示されます。カテゴリで最も時間のかかる XLA 演算の上位 10 個の表示と非表示を切り替えるには、テーブル内のカテゴリ名の横にある三角形をクリックします。

イメージ

  1. 時間: そのカテゴリに含まれるすべての演算の実行時間の合計を割合で表します。クリックするとエントリが展開され、個々の演算によって費やされた時間の内訳が表示されます。
  2. 上位 10 個の演算: カテゴリ名の横にあるトグルを使用すると、カテゴリ内で時間のかかる上位 10 個の演算を表示または非表示にできます。演算リストに fusion 演算が表示されている場合、このエントリを展開すると、fusion 以外の要素ごとの演算が表示されます。
  3. TensorFlow 演算: XLA 演算に関連付けられる TensorFlow 演算名が表示されます。
  4. FLOPS: FLOPS 使用率が表示されます。FLOPS の測定値が Cloud TPU ピーク FLOPS の割合として表示されます。FLOPS 使用率が高いほど、演算が速くなります。テーブルセルは色分けされています。FLOPS 使用率が高い場合には緑色(良好)、FLOPS 使用率が低い場合には赤色(不良)です。

演算詳細カード

テーブル エントリを選択すると、XLA 演算または演算カテゴリの詳細を表示するカードが表示されます。通常、次のようなカードが表示されます。

イメージ

  • 名前カテゴリ: ハイライト表示された XLA 演算名とカテゴリが表示されます。
  • FLOPS 使用率: FLOPS の使用率が FLOPS の合計数との割合で表示されます。
  • : 演算を含む XLA 式が表示されます。
  • メモリ使用率: プログラムによるピーク時のメモリ使用率が表示されます。
  • レイアウト(畳み込み演算のみ): XLA コンパイラによって実行されるパディングの説明を含む、テンソルの形状とレイアウトが表示されます。

結果の解釈

畳み込み演算の場合、TPU FLOPS の使用率が低い理由は、次のいずれかまたは両方です。

  • パディング(マトリックス ユニットが部分的に使用されている)
  • 畳み込み演算がメモリバインドになっている

このセクションでは、FLOP 使用率が低いモデルのパフォーマンス指標を解釈します。この例では、出力フュージョン畳み込み が実行時間の大半を占めています。FLOP 使用率が低いベクトルやスカラー オペレーションが数多く存在していました。

このようなプロファイルを最適化する方法としては、ベクトル演算またはスカラー演算を畳み込み演算に変換する方法があります。

次の例で、%convolution.399 は前の例の %convolution.340 よりも FLOPS とメモリ使用率が低くなっています。

イメージ

この例では、バッチサイズは 128 にパディングされ、特徴サイズは 8 にパディングされます。この場合、マトリックス ユニットの 5% しか効率的に使用されていない。使用率は、((batch_time * num_of_features) / padding_size ) / num_of_cores) で計算されます。この例の FLOPS をパディングを使用しない前の例の %convolution.340 と比較します。

トレース ビューア

トレース ビューアは、[プロファイル] ページで利用可能な Cloud TPU パフォーマンス分析ツールです。このツールは、Chrome トレース イベント プロファイリング ビューアを使用するため、Chrome ブラウザでのみ機能します。

トレース ビューアには次のタイムラインが表示されます。

  • TensorFlow モデルによって実行された演算の実行期間。
  • 演算を実行したシステムの部分(TPU またはホストマシン)。通常、ホストマシンがトレーニング データを前処理して TPU に転送する infeed 演算を実行し、TPU は実際のモデル トレーニングを行います。

トレース ビューアを使用して、モデル内のパフォーマンスの問題を特定し、この問題を解決する対策を講じることができます。たとえば、大まかには、インフィードとモデル トレーニングのどちらに大部分の時間を費やしているかどうかを識別できます。さらに詳しく見ると、実行に最も時間がかかっている TensorFlow 演算も識別できます。

トレース ビューアは、Cloud TPU ごとに 100 万件のイベントに制限されています。より多くのイベントを評価する必要がある場合は、代わりにストリーミング トレース ビューアを使用してください。

トレース ビューアのインターフェース

トレース ビューアを開くには、TensorBoard に移動して画面上部の [プロファイル] タブをクリックし、[ツール] プルダウンから [trace_viewer] を選択します。ビューアが開き、最新の実行結果が表示されます。

イメージ

この画面には、次の主要な要素(前述のスクリーンショットで番号が付けられています)が含まれています。

  1. 実行プルダウン: トレース情報をキャプチャしたすべての実行が表示されます。デフォルトのビューは最新の実行ですが、プルダウンを開いて別の実行を選択することもできます。
  2. ツール プルダウン: さまざまなプロファイリング ツールを選択します。
  3. ホスト プルダウン: Cloud TPU セットを含むホストを選択します。
  4. [Timeline] ペイン: Cloud TPU とホストマシンで実行された演算が時系列で表示されます。
  5. 詳細ウィンドウペイン: [Timeline] ペインで選択した演算の詳細情報が表示されます。

では、[Timeline] ペインについて詳しく見てみましょう。

イメージ

[Timeline] ペインには、次の要素が含まれます。

  1. 上部バー: さまざまな補助コントロールが表示されます。
  2. 時間軸: トレースの開始位置を基準にした時間が表示されます。
  3. セクションとトラックラベル: 各セクションには複数のトラックが含まれています。左側にある三角形をクリックすると、セクションの展開や折りたたみを行うことができます。システムで処理中の要素ごとに 1 つのセクションがあります。
  4. ツールセレクタ: トレース ビューアを操作するさまざまなツールが用意されています。
  5. イベント: 演算が実行されていた時間やトレーニング ステップなどのメタイベントの期間を示します。
  6. 垂直タブバー: Cloud TPU での使用は適しません。このバーは、Chrome が提供する汎用のトレース ビューアツールの一部で、さまざまなパフォーマンス分析タスクに使用されます。

セクションとトラック

トレース ビューアには、次のセクションがあります。

  • TPU ノードごとに 1 つのセクション。ラベルとして TPU チップの数とチップ内の TPU ノードの数が使用されます(例: 「Chip 2: TPU Core 1」)。TPU ノードのセクションには、次のトラックが含まれます。
    • ステップ: TPU で実行されていたトレーニング ステップの期間が表示されます。
    • TensorFlow Ops:: TPU 上で実行される TensorFlow 演算。
    • XLA Ops: TPU 上で実行された XLA 演算が表示されます。1 つの演算が 1 つ以上の XLA 演算に変換されます。XLA コンパイラにより、XLA 演算が TPU 上で実行されるコードに変換されます。
  • ホスト CPU 上で実行されるスレッドのセクション。「Host Threads」というラベルが付いています。このセクションには、CPU スレッドごとに 1 つのトラックが含まれます。注: セクション ラベルと一緒に表示される情報は無視してもかまいません。

タイムラインのツールセレクタ

TensorBoard のタイムライン ツールセレクタを使用して、タイムライン ビューを操作できます。タイムライン ツールをクリックするか、次のキーボード ショートカットを使用してツールを有効にしてハイライトできます。タイムライン ツールセレクタを移動するには、上部の点線部分をクリックしてセレクタを目的の場所にドラッグします。

タイムライン ツールの使い方は次のとおりです。

選択ツール
イベントをクリックして選択します。または、複数のイベントを選択するにはドラッグします。選択したイベントに関する詳細情報(名前、開始時間、期間)が、詳細ペインに表示されます。

パンツール
ドラッグしてタイムライン ビューを水平方向または垂直方向にパンします。

ズームツール
水平方向(時間軸)に沿って上にドラッグするとズームインし、下にドラッグするとズームアウトします。マウスカーソルの水平位置により、ズームが行われる中心が決まります。

注: マウスボタンを離した後、ズームツールがアクティブのままの場合は、タイムライン ビューをクリックしてズームツールを無効にします。

タイミング ツール
水平方向にドラッグして、時間間隔をマークできます。間隔の長さは時間軸に表示されます。間隔を調整するには、間隔の両端をドラッグします。間隔をクリアするには、タイムライン ビュー内の任意の場所をクリックします。

別のツールを選択しても、間隔はマークされたままになります。

メモリビューア

メモリビューアには、プログラムのピーク時のメモリ使用量とメモリ使用量の傾向が視覚的に表示されます。

メモリビューアのユーザー インターフェースは、次のようになります。

イメージ

  1. [ホスト] プルダウン: 表示する TPU ホストと XLA High Level Optimizer(HLO)モジュールを選択します。
  2. メモリの概要: ピーク時のメモリ割り当てとパディングなしのサイズが表示されます。
  3. 作業スペースグラフ: ピーク時のメモリ使用量とプログラムのメモリ使用率の傾向が表示されます。バッファグラフに表示されているバッファにカーソルを合わせると、バッファ割り当てカードに追加情報が表示されます。
  4. バッファグラフ: ピーク時のメモリ使用量でのバッファ割り当てを表示する 2 つのグラフが表示されます。いずれかのバッファグラフに表示されているバッファにカーソルを合わせると、バッファの詳細カードに追加情報が表示されます。
  5. バッファ割り当て詳細カード: バッファの割り当て詳細が表示されます。

メモリ概要パネル

メモリの概要パネルには、合計バッファ割り当てサイズが最大に達したときのモジュール名とピークメモリ割り当てが表示されます。比較のため、パディングされていないピーク割り当てサイズも表示されます。

イメージ

作業スペースグラフ

このグラフには、ピーク時のメモリ使用量とプログラムのメモリ使用の傾向が表示されます。縦線は、プログラムのピーク時のメモリ使用率を示します。このグラフは、プログラムが使用可能なグローバル メモリ空間に収まるかどうかを示しています。

イメージ

グラフの各ポイントは、XLA HLO プログラムにおける「プログラム ポイント」を表します。この線は、プログラムのメモリ使用量の経時的な変化を示しています。

バッファグラフ要素の操作

バッファグラフでバッファにカーソルを合わせると、バッファの存続期間を示す水平線が作業スペースグラフに表示されます。

イメージ

水平線の太さは、ピーク時のメモリ割り当てに対する相対的なバッファサイズを表します。線の長さはバッファの存続期間を示します。

バッファグラフ

2 つのグラフで、ピーク時のメモリ使用量の内訳が表示されます。

イメージ

  • By Program Order: プログラムの実行中にアクティブだったバッファが左から順に表示されます。

  • By Size: プログラムの実行中にアクティブだったバッファが、サイズの降順で表示されます。

バッファ割り当て詳細カード

バッファグラフに表示されているバッファにカーソルを合わせると、バッファ割り当て詳細カードが表示されます。通常、次のような詳細カードが表示されます。

イメージ

  1. 名前 - XLA オペレーションの名前。
  2. カテゴリ - オペレーションのカテゴリ。
  3. サイズ - バッファ割り当てのサイズ(パディングを含む)。
  4. パディングなしのサイズ - パディングなしのバッファ割り当てサイズ。
  5. 拡張 - パディングありとパディングなしの場合のバッファサイズの違い。
  6. 追加メモリ - パディングのために余分に使用されるメモリ量を示します。
  7. シェイプ - N 次元配列のランク、サイズ、データ型。
  8. TensorFlow 演算名 - バッファ割り当てに関連した TensorFlow 演算の名前が表示されます。
  9. 割り当てタイプ: バッファ割り当てカテゴリ(パラメータ、出力、スレッド ローカル、一時(fusion 内のバッファ割り当てなど))を示します。

メモリ不足エラー

モデルの実行中にメモリ不足エラーが発生した場合は、このドキュメントのガイドラインを使用してプロファイルを取得します。スクリプトがモデルをトレーニングするまで待ってから、プロファイラを起動します。プロファイリングの出力は、エラーの原因を特定するのに役立ちます。