Cloud TPU v5e 推論コンバータの概要

はじめに

Cloud TPU Inference Converer は、TPU 推論用の TensorFlow 2（TF2）モデルを準備して最適化します。コンバータは、ローカルまたは TPU VM シェルで動作します。 コンバータに必要なコマンドライン ツールがプリインストールされているため、TPU VM シェルをおすすめします。エクスポートされた SavedModel を受け取り、次の手順を実行します。

1. TPU 変換: TPUPartitionedCall と他の TPU 演算をモデルに追加して、TPU で使用できるようにします。デフォルトでは、推論用にエクスポートされたモデルにはそのような演算がないため、TPU でトレーニングされた場合でも、TPU では提供できません。
2. バッチ処理: グラフ内のバッチ処理を有効にして、スループットを向上させるモデルにバッチ処理演算を追加します。
3. BFloat16 変換: モデルのデータ形式を float32 から bfloat16 に変換して、計算パフォーマンスを向上させ、TPU の高帯域幅メモリ（HBM）の使用量を低減します。
4. IO シェイプの最適化: CPU と TPU の間で転送されるデータのテンソル形状を最適化して、帯域幅使用率を向上させます。

モデルをエクスポートするとき、ユーザーは TPU で実行する関数の関数エイリアスを作成します。これらの関数を Converter に渡して、Converter が TPU に配置して最適化します。

Cloud TPU Inference Converter は、Docker がインストールされている任意の環境で実行できる Docker イメージとして利用できます。

上記の手順を完了するまでの推定所要時間: 約 20 分～ 30 分

前提条件

1. モデルは TF2 モデルで、SavedModel 形式でエクスポートする必要があります。
2. モデルには、TPU 関数の関数エイリアスが必要です。これを行う方法については、サンプルコードをご覧ください。次の例では、tpu_func を TPU 関数のエイリアスとして使用します。
3. Tensorflow ライブラリ（Cloud TPU Inference Converter の依存関係）は AVX 命令を使用するようにコンパイルされるため、マシンの CPU が Advanced Vector eXtensions（AVX）命令をサポートしていることを確認してください。ほとんどの CPU が AVX をサポートします。
   1. lscpu | grep avx を実行して、AVX 命令セットがサポートされているかどうかを確認できます。

準備

設定を開始する前に、次のことを行います。

• 新しいプロジェクトを作成する: Google Cloud コンソールの [プロジェクト セレクタ] ページで、Cloud プロジェクトを選択または作成します。

• TPU VM を設定する: Google Cloud Console または gcloud を使用して新しい TPU VM を作成するか、既存の TPU VM を使用して変換されたモデルで推論を実行します。

  • TPU VM イメージが TensorFlow ベースであることを確認します。例: --version=tpu-vm-tf-2.11.0
  • 変換されたモデルがこの TPU VM に読み込まれ、提供されます。

• Cloud TPU Inference Converter を使用するために必要なコマンドライン ツールがあることを確認します。Google Cloud SDK と Docker はローカルにインストールすることも、このソフトウェアがデフォルトでインストールされている TPU VM を使用することもできます。これらのツールを使用して、Converter のイメージを操作します。

  次のコマンドを使用して、SSH を使用してインスタンスに接続します。

  gcloud compute tpus tpu-vm ssh ${tpu-name} --zone ${zone} --project ${project-id}
  

環境設定

TPU VM シェルまたはローカルシェルから環境を設定します。

TPU VM シェル

• TPU VM シェルで次のコマンドを実行して、root 以外の Docker の使用を許可します。

  sudo usermod -a -G docker ${USER}
  newgrp docker
  

• Docker 認証ヘルパーを初期化します。

  gcloud auth configure-docker \
    us-docker.pkg.dev
  

ローカルシェル

ローカルシェルで、次の手順で環境を設定します。

• Cloud SDK をインストールします。これには、gcloud コマンドライン ツールが含まれています。

• Docker をインストールします。

• root 以外の Docker 使用を許可します。

  sudo usermod -a -G docker ${USER}
  newgrp docker
  

• 環境にログインします。

  gcloud auth login
  

• Docker 認証ヘルパーを初期化します。

  gcloud auth configure-docker \
      us-docker.pkg.dev
  

• Inference Converter の Docker イメージを pull します。

    CONVERTER_IMAGE=us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
    docker pull ${CONVERTER_IMAGE}
    

コンバータのイメージ

イメージは 1 回限りのモデル変換用です。モデルパスを設定し、必要に応じてコンバータ オプションを調整します。使用例セクションには、いくつかの一般的なユースケースが用意されています。

注: 次のコードでは、docker コマンドで mount を指定せずに、Cloud Storage バケットを input_model_dir と output_model_dir に使用できます。

docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      function_alias: "tpu_func"
    }
    batch_options {
      num_batch_threads: 2
      max_batch_size: 8
      batch_timeout_micros: 5000
      allowed_batch_sizes: 2
      allowed_batch_sizes: 4
      allowed_batch_sizes: 8
      max_enqueued_batches: 10
    }
'

TPU VM で変換されたモデルによる推論

# Initialize the TPU
resolver = tf.distribute.cluster_resolver.TPUClusterResolver("local")
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

# Load the model
model = tf.saved_model.load(${CONVERTED_MODEL_PATH})

# Find the signature function for serving
serving_signature = 'serving_default' # Change the serving signature if needed
serving_fn = model.signatures[serving_signature]
# Run the inference using requests.
results = serving_fn(**inputs)
logging.info("Serving results: %s", str(results))

使用例

TPU 関数の関数エイリアスを追加する

1. TPU で実行するすべての関数をラップする関数をモデル内で検索または作成します。@tf.function が存在しない場合は追加します。
2. モデルを保存するときに、次のような SaveOptions を指定して、model.tpu_func にエイリアス func_on_tpu を付与します。
3. この関数のエイリアスをコンバータに渡すことができます。

class ToyModel(tf.keras.Model):
  @tf.function(
      input_signature=[tf.TensorSpec(shape=[None, 10], dtype=tf.float32)])
  def tpu_func(self, x):
    return x * 1.0

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'func_on_tpu': model.tpu_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

複数の TPU 関数でモデルを変換する

TPU には複数の関数を配置できます。複数の関数エイリアスを作成し、converter_options_string でコンバータに渡します。

tpu_functions {
  function_alias: "tpu_func_1"
}
tpu_functions {
  function_alias: "tpu_func_2"
}

量子化

重要: この機能は試験運用中であり、予告なく変更されることがあります。量子モデルのパフォーマンスを本番環境で使用する前に、パフォーマンスを慎重に評価する必要があります。量子化の TPU v5e のサポートは、TensorFlow で提供されています。

量子化は、モデルのパラメータを表すために使用される数値の精度を下げる手法です。この結果、モデルサイズが小さくなり、計算速度が向上します。量子化モデルでは、推論のスループットが向上するだけでなく、メモリ使用量とストレージ サイズも小さくなりますが、そのかわりに精度は低くなります。

TPU を対象とする TensorFlow の新しいトレーニング後の量子化機能は、モバイル デバイスとエッジデバイスを対象とする TensorFlow Lite の同様の既存の機能から開発されています。一般的な量子化の詳細については、TensorFlow Lite のドキュメントをご覧ください。

量子化のコンセプト

このセクションでは、推論コンバータを使用した量子化に関連するコンセプトを定義します。

他の TPU 構成（スライス、ホスト、チップ、TensorCore など）に関連するコンセプトについては、TPU システム アーキテクチャのページをご覧ください。

• トレーニング後の量子化（PTQ）: PTQ は、精度を大幅に下げることなく、ニューラル ネットワーク モデルのサイズとコンピューティングの複雑さを下げる方法です。PTQ は、トレーニング済みモデルの浮動小数点重みと活性化を 8 ビット整数や 16 ビット整数などの低精度整数に変換することで機能します。これにより、モデルのサイズと推論レイテンシが大幅に低下する一方で、精度はわずかに低下します。

• キャリブレーション: 量子化のキャリブレーション ステップは、ニューラル ネットワーク モデルの重みと活性化が取る値の範囲に関する統計情報を収集するプロセスです。この情報は、モデルの量子化パラメータを決定するために使用されます。これは浮動小数点の重みと活性化を整数に変換するために使用される値です。

• 代表データセット: 量子化の代表データセットは、モデルの実際の入力データを表す小さなデータセットです。量子化のキャリブレーション ステップで、モデルの重みと活性化が取る値の範囲に関する統計情報を収集します。代表的なデータセットは次のプロパティを満たす必要があります。

  • 推論中のモデルへの実際の入力を適切に表す必要があります。つまり、モデルが現実に認識する可能性のある値の範囲をカバーする必要があります。
  • 存在する場合は、条件の各ブランチ（tf.cond など）を一緒に通過する必要があります。これは重要です。量子化プロセスでは、代表的なデータセットで明示的に表現されていない場合でも、モデルへのすべての入力を処理できる必要があります。
  • 十分な統計情報を収集し、エラーを減らすために十分な大きさである必要があります。目安としては、200 を超える代表サンプルを使用することをおすすめします。

  代表データセットは、トレーニング データセットのサブセットにすることも、モデルへの実際の入力を表すように特別に設計された別のデータセットにすることもできます。使用するデータセットの選択は、アプリケーションによって異なります。

• 静的範囲量子化（SRQ）: SRQ では、キャリブレーション ステップ中にニューラル ネットワーク モデルの重みと活性化の値の範囲を 1 回決定します。つまり、モデルへのすべての入力に同じ値の範囲が使用されます。これは、特に広い範囲の入力値を持つモデルの場合、動的範囲の量子化よりも精度が低くなります。ただし、静的範囲量子化では、動的範囲の量子化よりも実行時の計算が少なくなります。

• 動的範囲量子化（DRQ）: DRQ は、各入力のニューラル ネットワーク モデルの重みと活性化の値の範囲を決定します。これにより、入力データの値の範囲にモデルを適応させ、精度を改善できます。ただし、動的範囲量子化には静的範囲量子化よりも実行時の多くの計算が必要です。

  機能	静的範囲量子化	動的範囲量子化
  値の範囲	キャリブレーション中に 1 回決定	入力ごとに決定
  精度	特に入力値の範囲が大きいモデルの場合は、精度が低下する可能性があります。	特に入力値の範囲が大きいモデルの場合は、より正確な値になる可能性があります。
  複雑さ	シンプルになった	複雑になった
  実行時の計算	少ない計算	多い計算

• 重みのみの量子化: 重みのみの量子化は、ニューラル ネットワーク モデルの重みのみを量子化して、活性化を浮動小数点のままにする量子化の種類です。モデルの精度を維持できるため、精度が重視されるモデルにおすすめのオプションです。

量子化の使用方法

量子化を適用するには、QuantizationOptions を構成してコンバータのオプションに設定します。次のようなオプションがあります。

• タグ: 量子化する SavedModel 内の MetaGraphDef を識別するタグのコレクション。MetaGraphDef が 1 つしかない場合、指定する必要はありません。
• signature_keys: 入力と出力を含む SignatureDef を識別するキー シーケンス。指定しない場合、[serving_default] が使用されます。
• quantization_method: 適用する量子化メソッド。指定しない場合は STATIC_RANGE 量子化が適用されます。
• op_set: XLA として保持する必要があります。これは現在、デフォルトのオプションであり、指定する必要はありません。
• representative_datasets: 量子化パラメータの調整に使用するデータセットを指定します。

代表的なデータセットを作成する

代表的なデータセットは、基本的にサンプルのイテラブルです。ここで、サンプルは {input_key: input_value} のマップです。例:

representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]

代表的なデータセットは、tf-nightly pip パッケージ内で現在利用可能な TfRecordRepresentativeDatasetSaver クラスを使用して、TFRecord ファイルとして保存する必要があります。例:

# Assumed tf-nightly installed.
import tensorflow as tf
representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]
tf.quantization.experimental.TfRecordRepresentativeDatasetSaver(
       path_map={'serving_default': '/tmp/representative_dataset_path'}
    ).save({'serving_default': representative_dataset})

例

次の例では、serving_default の署名キーと tpu_func の関数エイリアスを使用してモデルを量子化します。

注: 次のコードでは、CONVERTER_IMAGE を us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:nightly として設定し、夜間ビルドを使用します。

docker run \
  --mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
  --mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
  ${CONVERTER_IMAGE} \
  --input_model_dir=/tmp/input \
  --output_model_dir=/tmp/output \
  --converter_options_string=' \
    tpu_functions { \
      function_alias: "tpu_func" \
    } \
    external_feature_configs { \
      quantization_options { \
        signature_keys: "serving_default" \
        representative_datasets: { \
          key: "serving_default" \
          value: { \
            tfrecord_file_path: "${TF_RECORD_FILE}" \
          } \
        } \
      } \
    } '

バッチ処理を追加する

Converter を使用して、モデルにバッチ処理を追加できます。調整可能なバッチ オプションの説明については、バッチ処理オプションの定義をご覧ください。

デフォルトでは、Converter はモデル内のすべての TPU 関数をバッチ処理します。また、ユーザー提供の署名と関数もバッチ処理できます。これにより、パフォーマンスをさらに向上させることができます。バッチ化される TPU 関数、ユーザー指定の関数、または署名は、バッチ処理演算の厳密な形状の要件を満たしている必要があります。

Converter は、既存のバッチ処理オプションを更新することもできます。 以下に、モデルにバッチ処理を追加する方法を示します。バッチ処理の詳細については、バッチ処理の詳細をご覧ください。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

bfloat16 と IO の形状の最適化を無効にする

BFloat16 と IO 形状の最適化はデフォルトで有効になっています。それらがモデルで適切に機能しない場合は、無効にできます。

# Disable both optimizations
disable_default_optimizations: true

# Or disable them individually
io_shape_optimization: DISABLED
bfloat16_optimization: DISABLED

変換レポート

この変換レポートは、推論コンバータの実行後にログで確認できます。下記の例をご覧ください。

-------- Conversion Report --------
TPU cost of the model: 96.67% (2034/2104)
CPU cost of the model:  3.33% (70/2104)

Cost breakdown
================================
%         Cost    Name
--------------------------------
3.33      70      [CPU cost]
48.34     1017    tpu_func_1
48.34     1017    tpu_func_2
--------------------------------

このレポートでは、CPU と TPU での出力モデルの計算コストを見積もり、さらに TPU の費用を各関数に分解します。これには、コンバータ オプションで TPU 関数の選択が反映されます。

TPU をさらに活用するには、モデル構造を試してコンバータ オプションを調整することをおすすめします。

注: 現在の変換レポートでは、TPU 関数のすべての演算は TPU 費用としてカウントされます。

よくある質問

TPU にはどの関数を配置すればよいですか。

大半の演算は TPU で高速に実行されるため、できるだけ多くのモデルを TPU に配置することをおすすめします。

モデルに TPU と互換性のない演算、文字列、スパース テンソルが含まれていない場合は、通常、モデル全体を TPU に配置するのが最適です。これを行うには、モデル全体をラップする関数を見つけるか作成し、その関数の関数エイリアスを作成して、Converter に渡します。

TPU で動作できない部分（TPU と互換性のない演算、文字列、スパース テンソルなど）がモデルに含まれる場合、TPU 関数は、互換性のない部分がどこにあるかによって異なります。

• 互換性のない部分がモデルの先頭または末尾にある場合、モデルをリファクタリングして CPU で維持できます。たとえば、文字列の前処理と後処理のステージなどです。コードを CPU に移動する方法について詳しくは、「モデルの一部を CPU に移動するにはどうすればよいですか？」をご覧ください。 モデルをリファクタリングする一般的な方法を示しています。
• モデルの中間に位置する場合、モデルを 3 つの部分に分割し、中間部分に TPU 互換でない演算をすべて含めて、CPU で実行することをおすすめします。
• スパース テンソルの場合は、CPU で tf.sparse.to_dense を呼び出し、結果のスパース テンソルをモデルの TPU 部分に渡すことを検討してください。

考慮すべきもう 1 つの要素は、HBM 使用量です。テーブルを埋め込むと、大量の HBM を使用することがあります。TPU のハードウェア制限を超えて増加した場合は、ルックアップ演算とともに CPU に配置する必要があります。

可能であれば、1 つのシグネチャ内に 1 つの TPU 関数のみを配置してください。 モデルの構造で、受信する推論リクエストごとに複数の TPU 関数を呼び出す必要がある場合は、CPU と TPU の間でテンソルを送信する場合のレイテンシが増加することに注意してください。

TPU 関数の選択を評価するには、変換レポートを確認することをおすすめします。TPU に配置された計算の割合と、各 TPU 関数の費用の内訳が表示されます。

モデルの一部を CPU に移動するにはどうすればよいですか？

TPU で提供できない部分がモデルに含まれている場合は、モデルをリファクタリングして CPU に移動する必要があります。玩具のサンプルを示します。このモデルは、前処理ステージがある言語モデルです。わかりやすくするために、レイヤの定義と関数のコードは省略しています。

class LanguageModel(tf.keras.Model):
  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.bert_layer(word_ids)

このモデルは、2 つの理由で TPU で直接提供できません。まず、パラメータが文字列です。次に、preprocess 関数には多数の文字列演算を含めることができます。どちらも TPU と互換性がありません。

このモデルをリファクタリングするには、tpu_func という別の関数を作成して、計算量の多い bert_layer をホストします。次に、tpu_func の関数エイリアスを作成し、Converter に渡します。このようにして、tpu_func 内のすべてのものが TPU 上で実行され、model_func に残されたすべてのものが CPU 上で実行されます。

class LanguageModel(tf.keras.Model):
  @tf.function
  def tpu_func(self, word_ids):
    return self.bert_layer(word_ids)

  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.tpu_func(word_ids)

モデルに TPU と互換性のない演算、文字列、スパース テンソルがある場合は、どうすればよいですか？

TPU では標準の TensorFlow 演算のほとんどがサポートされていますが、スパース テンソルや文字列などはサポートされていません。Converter は TPU と互換性のないオペレーションを確認しません。したがって、このような演算を含むモデルは変換をパスすることができます。しかし、推論に実行すると、次のようなエラーが発生します。

'tf.StringToNumber' op isn't compilable for TPU device.

モデルに TPU と互換性のない演算がある場合は、TPU 関数の外部に配置する必要があります。さらに、文字列は TPU でサポートされていないデータ形式です。したがって、文字列型変数は TPU 関数に配置しないでください。また、TPU 関数のパラメータと戻り値も文字列型にしないでください。同様に、TPU 関数には、パラメータや戻り値を含め、スパース テンソルを配置しないでください。

通常、モデルの互換性のない部分をリファクタリングして CPU に移動することは困難ではありません。ここで例を示します。

モデルでカスタム演算をサポートする方法

モデルでカスタム演算が使用されている場合、Converter はモデルを認識せず、モデルの変換に失敗する場合があります。これは、演算の完全な定義を含むカスタム演算の演算ライブラリが Converter にリンクされていないためです。

コンバータのコードはまだオープンソースではないため、コンバータはカスタム演算でビルドできません。

TensorFlow 1 モデルを使用している場合はどうすればよいですか？

Converter は TensorFlow 1 モデルをサポートしていません。TensorFlow 1 モデルを TensorFlow 2 に移行する必要があります。

モデルを実行するときに MLIR ブリッジを有効にする必要がありますか？

変換されたほとんどのモデルは、新しい TF2XLA MLIR ブリッジまたは元の TF2XLA ブリッジのいずれかで実行できます。

関数エイリアスを使用せずに、すでにエクスポートされているモデルを変換するにはどうすればよいですか？

関数エイリアスを使用せずにモデルがエクスポートされた場合は、そのモデルを再度エクスポートして、関数エイリアスを作成するのが最も簡単な方法です。 再エクスポートができない場合でも、concrete_function_name を指定してモデルを変換できます。ただし、正しい concrete_function_name を識別するには、いくつかの検出作業が必要です。

関数エイリアスは、ユーザー定義の文字列から具体的な関数名へのマッピングです。モデル内の特定の関数を参照しやすくなります。Converter は、関数エイリアスと未加工の具体的な関数名の両方を受け入れます。

具体的な関数名は saved_model.pb を調べることで確認できます。

次の例は、__inference_serve_24 という具体的な関数を TPU に配置する方法を示しています。

sudo docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      concrete_function_name: "__inference_serve_24"
    }'

コンパイル時の定数制約エラーを解決するにはどうすればよいですか？

トレーニングと推論の両方で、XLA では TPU のコンパイル時に、特定の演算への入力が既知の形状である必要があります。つまり、XLA がプログラムの TPU 部分をコンパイルする場合、これらの演算への入力は静的に既知の形状でなければなりません。

この問題を解決するには次の 2 つの方法があります。

• 最良のオプションは、XLA が TPU プログラムをコンパイルするまでに、演算の入力を更新して、静的に既知の形状を維持することです。このコンパイルは、モデルの TPU 部分が実行される直前に行われます。つまり、TpuFunction が実行されるときには、形状が静的に既知である必要があります。
• もう 1 つの方法は、問題のある演算を含めないように TpuFunction を変更することです。

バッチ処理の形状のエラーが表示されるのはなぜですか？

バッチ処理には、受信リクエストを 0 番目のディメンション（バッチ処理ディメンション）に沿ってバッチ化できる、厳密な形状の要件があります。これらの形状要件は TensorFlow バッチ処理演算に由来するため、緩和することはできません。

これらの要件を満たさないと、次のようなエラーが発生します。

1. バッチ処理の入力テンソルには少なくとも 1 つのディメンションが必要です。
2. 入力のディメンションは一致している必要があります。
3. 特定の演算呼び出しで指定されたバッチ入力テンソルは、同じ 0 次元のディメンションサイズである必要があります。
4. バッチ化される出力テンソルの 0 番目のディメンションサイズは、入力テンソルの 0 番目のディメンションサイズの合計と等しくありません。

これらの要件を満たすには、バッチに別の関数または署名を提供することを検討してください。これらの要件を満たすために既存の関数を変更する必要がある場合もあります。

関数がバッチ化されている場合は、@tf.function の input_signature の形状がすべて、0 番目のディメンションで None になっていることを確認します。署名がバッチ化されている場合は、その入力がすべて、0 番目のディメンションで -1 になっていることを確認します。

これらのエラーが発生する理由とその解決方法については、バッチ処理の詳細をご覧ください。

既知の問題

TPU 関数は別の TPU 関数を間接的に呼び出すことができない

Converter は CPU-TPU 境界を越えてほとんどの関数呼び出しシナリオを処理できますが、まれに失敗するエッジケースが 1 つあります。TPU 関数が別の TPU 関数を間接的に呼び出すときです。

これは、Converter が、TPU 関数自体の呼び出しから TPU 呼び出しスタブの呼び出しまで、TPU 関数の直接の呼び出し元を変更するためです。呼び出しスタブには、CPU でのみ動作する演算が含まれています。最終的に直接の呼び出し元を呼び出す関数を TPU 関数が呼び出すと、それらの CPU 演算が TPU で実行され、カーネル エラーが発生する可能性があります。このケースは、別の TPU 関数を直接呼び出す TPU 関数とは異なります。この場合、コンバータはどちらの関数もコールスタブを呼び出すように変更しないため、正常に動作します。

Converter では、このシナリオの検出を実装しました。次のエラーが表示された場合、モデルはこのエッジケースに到達しています。

Unable to place both "__inference_tpu_func_2_46" and "__inference_tpu_func_4_68"
on the TPU because "__inference_tpu_func_2_46" indirectly calls
"__inference_tpu_func_4_68". This behavior is unsupported because it can cause
invalid graphs to be generated.

一般的な解決策は、このような関数呼び出しのシナリオを回避するためにモデルをリファクタリングすることです。難しい場合は、Google サポートチームに連絡して詳細を確認してください。

リファレンス

Protobuf 形式のコンバータ オプション

message ConverterOptions {
  // TPU conversion options.
  repeated TpuFunction tpu_functions = 1;

  // The state of an optimization.
  enum State {
    // When state is set to default, the optimization will perform its
    // default behavior. For some optimizations this is disabled and for others
    // it is enabled. To check a specific optimization, read the optimization's
    // description.
    DEFAULT = 0;
    // Enabled.
    ENABLED = 1;
    // Disabled.
    DISABLED = 2;
  }

  // Batch options to apply to the TPU Subgraph.
  //
  // At the moment, only one batch option is supported. This field will be
  // expanded to support batching on a per function and/or per signature basis.
  //
  //
  // If not specified, no batching will be done.
  repeated BatchOptions batch_options = 100;

  // Global flag to disable all optimizations that are enabled by default.
  // When enabled, all optimizations that run by default are disabled. If a
  // default optimization is explicitly enabled, this flag will have no affect
  // on that optimization.
  //
  // This flag defaults to false.
  bool disable_default_optimizations = 202;

  // If enabled, apply an optimization that reshapes the tensors going into
  // and out of the TPU. This reshape operation improves performance by reducing
  // the transfer time to and from the TPU.
  //
  // This optimization is incompatible with input_shape_opt which is disabled.
  // by default. If input_shape_opt is enabled, this option should be
  // disabled.
  //
  // This optimization defaults to enabled.
  State io_shape_optimization = 200;

  // If enabled, apply an optimization that updates float variables and float
  // ops on the TPU to bfloat16. This optimization improves performance and
  // throughtput by reducing HBM usage and taking advantage of TPU support for
  // bfloat16.
  //
  // This optimization may cause a loss of accuracy for some models. If an
  // unacceptable loss of accuracy is detected, disable this optimization.
  //
  // This optimization defaults to enabled.
  State bfloat16_optimization = 201;

  BFloat16OptimizationOptions bfloat16_optimization_options = 203;

  // The settings for XLA sharding. If set, XLA sharding is enabled.
  XlaShardingOptions xla_sharding_options = 204;
}

message TpuFunction {
  // The function(s) that should be placed on the TPU. Only provide a given
  // function once. Duplicates will result in errors. For example, if
  // you provide a specific function using function_alias don't also provide the
  // same function via concrete_function_name or jit_compile_functions.
  oneof name {
    // The name of the function alias associated with the function that
    // should be placed on the TPU. Function aliases are created during model
    // export using the tf.saved_model.SaveOptions.
    //
    // This is a recommended way to specify which function should be placed
    // on the TPU.
    string function_alias = 1;

    // The name of the concrete function that should be placed on the TPU. This
    // is the name of the function as it found in the GraphDef and the
    // FunctionDefLibrary.
    //
    // This is NOT the recommended way to specify which function should be
    // placed on the TPU because concrete function names change every time a
    // model is exported.
    string concrete_function_name = 3;

    // The name of the signature to be placed on the TPU. The user must make
    // sure there is no TPU-incompatible op under the entire signature.
    string signature_name = 5;

    // When jit_compile_functions is set to True, all jit compiled functions
    // are placed on the TPU.
    //
    // To use this option, decorate the relevant function(s) with
    // @tf.function(jit_compile=True), before exporting. Then set this flag to
    // True. The converter will find all functions that were tagged with
    // jit_compile=True and place them on the TPU.
    //
    // When using this option, all other settings for the TpuFunction
    // will apply to all functions tagged with
    // jit_compile=True.
    //
    // This option will place all jit_compile=True functions on the TPU.
    // If only some jit_compile=True functions should be placed on the TPU,
    // use function_alias or concrete_function_name.
    bool jit_compile_functions = 4;
  }

}

message BatchOptions {
  // Number of scheduling threads for processing batches of work. Determines
  // the number of batches processed in parallel. This should be roughly in line
  // with the number of TPU cores available.
  int32 num_batch_threads = 1;

  // The maximum allowed batch size.
  int32 max_batch_size = 2;

  // Maximum number of microseconds to wait before outputting an incomplete
  // batch.
  int32 batch_timeout_micros = 3;

  // Optional list of allowed batch sizes. If left empty,
  // does nothing. Otherwise, supplies a list of batch sizes, causing the op
  // to pad batches up to one of those sizes. The entries must increase
  // monotonically, and the final entry must equal max_batch_size.
  repeated int32 allowed_batch_sizes = 4;

  // Maximum number of batches enqueued for processing before requests are
  // failed fast.
  int32 max_enqueued_batches = 5;

  // If set, disables large batch splitting which is an efficiency improvement
  // on batching to reduce padding inefficiency.
  bool disable_large_batch_splitting = 6;

  // Experimental features of batching. Everything inside is subject to change.
  message Experimental {
    // The component to be batched.
    // 1. Unset if it's for all TPU subgraphs.
    // 2. Set function_alias or concrete_function_name if it's for a function.
    // 3. Set signature_name if it's for a signature.
    oneof batch_component {
      // The function alias associated with the function. Function alias is
      // created during model export using the tf.saved_model.SaveOptions, and is
      // the recommended way to specify functions.
      string function_alias = 1;

      // The concreate name of the function. This is the name of the function as
      // it found in the GraphDef and the FunctionDefLibrary. This is NOT the
      // recommended way to specify functions, because concrete function names
      // change every time a model is exported.
      string concrete_function_name = 2;

      // The name of the signature.
      string signature_name = 3;
    }
  }

  Experimental experimental = 7;
}

message BFloat16OptimizationOptions {
  // Indicates where the BFloat16 optimization should be applied.
  enum Scope {
    // The scope currently defaults to TPU.
    DEFAULT = 0;
    // Apply the bfloat16 optimization to TPU computation.
    TPU = 1;
    // Apply the bfloat16 optimization to the entire model including CPU
    // computations.
    ALL = 2;
  }

  // This field indicates where the bfloat16 optimization should be applied.
  //
  // The scope defaults to TPU.
  Scope scope = 1;

  // If set, the normal safety checks are skipped. For example, if the model
  // already contains bfloat16 ops, the bfloat16 optimization will error because
  // pre-existing bfloat16 ops can cause issues with the optimization. By
  // setting this flag, the bfloat16 optimization will skip the check.
  //
  // This is an advanced feature and not recommended for almost all models.
  //
  // This flag is off by default.
  bool skip_safety_checks = 2;

  // Ops that should not be converted to bfloat16.
  // Inputs into these ops will be cast to float32, and outputs from these ops
  // will be cast back to bfloat16.
  repeated string filterlist = 3;
}

message XlaShardingOptions {
  // num_cores_per_replica for TPUReplicateMetadata.
  //
  // This is the number of cores you wish to split your model into using XLA
  // SPMD.
  int32 num_cores_per_replica = 1;

  // (optional) device_assignment for TPUReplicateMetadata.
  //
  // This is in a flattened [x, y, z, core] format (for
  // example, core 1 of the chip
  // located in 2,3,0 will be stored as [2,3,0,1]).
  //
  // If this is not specified, then the device assignments will utilize the same
  // topology as specified in the topology attribute.
  repeated int32 device_assignment = 2;

  // A serialized string of tensorflow.tpu.TopologyProto objects, used for
  // the topology attribute in TPUReplicateMetadata.
  //
  // You must specify the mesh_shape and device_coordinates attributes in
  // the topology object.
  //
  // This option is required for num_cores_per_replica > 1 cases due to
  // ambiguity of num_cores_per_replica, for example,
  // pf_1x2x1 with megacore and df_1x1
  // both have num_cores_per_replica = 2, but topology is (1,2,1,1) for pf and
  // (1,1,1,2) for df.
  // - For pf_1x2x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,2,1,1]
  //   device_coordinates=flatten([0,0,0,0], [0,1,0,0])
  // - For df_1x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,1,1,2]
  //   device_coordinates=flatten([0,0,0,0], [0,0,0,1])
  // - For df_2x2, mesh shape and device_coordinates looks like:
  //   mesh_shape = [2,2,1,2]
  //   device_coordinates=flatten(
  //    [0,0,0,0],[0,0,0,1],[0,1,0,0],[0,1,0,1]
  //    [1,0,0,0],[1,0,0,1],[1,1,0,0],[1,1,0,1])
  bytes topology = 3;
}

バッチ処理の詳細

バッチ処理は、スループットと TPU の使用率の向上に使用されます。これにより、複数のリクエストを同時に処理できます。トレーニング中、tf.data を使用してバッチ処理を実行できます。推論の際には、通常、受信リクエストをバッチ処理する演算をグラフに追加します。この演算は、十分なリクエストが発生するか、タイムアウトに到達して個々のリクエストから大きなバッチを生成するまで待機します。バッチサイズやタイムアウトなど、調整できるさまざまなバッチ処理オプションの詳細については、バッチ処理オプションの定義をご覧ください。

デフォルトでは、Converter は TPU の計算の直前にバッチ処理演算を挿入します。ユーザー指定の TPU 関数と、モデル内の既存の TPU 計算をバッチ処理演算でラップします。このデフォルトの動作は、コンバータにどの関数や署名をバッチ化するかを指示することでオーバーライドできます。

次の例は、デフォルトのバッチ処理を追加する方法を示しています。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

署名のバッチ処理

注: この機能はまだ試験運用版です。テスト後、このフィールドは試験運用版のラッパーから移動されます。

署名のバッチ処理では、モデルの入力全体を署名の入力から署名の出力へとバッチ処理します。Converter のデフォルトのバッチ処理動作とは異なり、署名のバッチ処理では、TPU の計算と CPU の計算の両方がバッチ処理されます。これにより、一部のモデルでの推論中に 10% ～ 20% のパフォーマンスが向上します。

すべてのバッチ処理と同様に、署名のバッチ処理には厳密な形状の要件があります。 これらの形状要件を満たすには、署名入力に 2 つ以上のディメンションを持つ形状が必要です。最初のディメンションはバッチサイズで、サイズは -1 になります。たとえば、(-1, 4)、(-1)、(-1, 128, 4, 10) はすべて有効な入力形状です。これが可能でない場合は、デフォルトのバッチ処理または関数のバッチ処理の使用を検討してください。

署名のバッチ処理を使用するには、BatchOptions を使用して署名名を signature_name として指定します。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    signature_name: "serving_default"
  }
}

関数のバッチ処理

注: この機能はまだ試験運用版です。テスト後、このフィールドは試験運用版のラッパーから移動されます。

関数のバッチ処理は、バッチ化する関数を Converter に指示するために使用できます。デフォルトでは、Converter はすべての TPU 関数をバッチ処理します。関数のバッチ処理は、このデフォルトの動作をオーバーライドします。

関数のバッチ処理は、CPU 計算のバッチ処理に使用できます。多くのモデルでは、CPU 計算がバッチ化されるとパフォーマンスが向上します。CPU 計算をバッチ処理するには、署名のバッチ処理を使用しますが、一部のモデルでは動作しない場合があります。このような場合、TPU 計算に加えて、CPU 計算の部分をバッチ処理するために関数のバッチ処理を使用できます。バッチ処理演算は TPU では実行できないため、提供されたバッチ処理関数は CPU で呼び出す必要があります。

関数のバッチ処理は、バッチ処理演算によって課される厳密な形状の要件を満たすためにも使用できます。TPU 関数がバッチ処理演算の形状要件を満たさない場合、関数のバッチ処理を使用して、異なる関数をバッチ処理するよう Converter に指示できます。

これを使用するには、バッチ化する関数の function_alias を生成します。これを行うには、バッチ化するすべてのものをラップする関数をモデル内で検出または作成します。この関数が、バッチ処理オプションによって課される厳密な形状の要件を満たしていることを確認します。@tf.function をまだ追加していない場合は、追加します。 input_signature を @tf.function に渡すことが重要です。0 番目のディメンションは、None である必要があります。これはバッチ ディメンションであり、固定サイズにできないためです。たとえば、[None, 4]、[None]、[None, 128, 4, 10] はすべて有効な入力シェイプです。モデルを保存するときに、次のような SaveOptions を指定して、model.batch_func にエイリアス「batch_func」を割り当てます。次に、この関数エイリアスをコンバータに渡すことができます。

class ToyModel(tf.keras.Model):
  @tf.function(input_signature=[tf.TensorSpec(shape=[None, 10],
                                              dtype=tf.float32)])
  def batch_func(self, x):
    return x * 1.0

  ...

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'batch_func': model.batch_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

次に、BatchOptions を使用して function_alias を渡します。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    function_alias: "batch_func"
  }
}

バッチ処理オプションの定義

• num_batch_threads:（整数）作業のバッチを処理するスケジューリング スレッドの数。並列に処理されるバッチの数を決定します。これは利用可能な TPU コアの数とほぼ同じです。
• max_batch_size:（整数）最大許容バッチサイズ。大規模なバッチ分割を行うには、allowed_batch_sizes より大きくできます。
• batch_timeout_micros:（整数）不完全なバッチを出力するまで待機する最大時間（マイクロ秒単位）。
• allowed_batch_sizes:（整数のリスト）リストが空でない場合、リスト内の最も近いサイズにバッチがパディングされます。リストは単調に増加し、最終的な要素は max_batch_size 以下でなければなりません。
• max_enqueued_batches:（整数）リクエストが失敗する前に処理するためにキューに格納されたバッチの最大数。

既存のバッチ処理オプションの更新

batch_options を指定して Docker イメージを実行し、--converter_options_string フラグを使用して disable_default_optimizations を true に設定すると、バッチ処理オプションを追加または更新できます。バッチ オプションは、すべての TPU 関数または既存のバッチ処理演算に適用されます。

注: --converter_options_string を使用して tpu_functions を一緒に設定して、バッチ オプションを更新することはできません。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}
disable_default_optimizations=True

バッチ処理の形状要件

バッチは、バッチ（0 番目）ディメンションに沿ってリクエスト間で入力テンソルを連結することによって作成されます。出力テンソルは 0 番目のディメンションに沿って分割されます。これらの演算を実行するために、バッチ処理ではその入力と出力について厳密な形状要件があります。

チュートリアル

これらの要件を理解するには、最初にバッチ処理がどのように行われるかを理解しておくことが役立ちます。以下の例では、単純な tf.matmul 演算をバッチ処理しています。

def my_func(A, B)
    return tf.matmul(A, B)

最初の推論リクエストは、(1, 3, 2) と (1, 2, 4) の形状で入力 A と B を生成します。2 番目の推論リクエストは、(2, 3, 2) と (2, 2, 4) の形状で入力 A と B を生成します。

バッチ処理のタイムアウトに到達しました。このモデルはバッチサイズ 3 をサポートしているため、推論リクエスト #1 と #2 はパディングなしで一緒にバッチ化されます。バッチ テンソルは、バッチ（0 番目）ディメンションに沿ってリクエスト #1 と #2 を連結することで形成されます。#1 の A の形状は (1, 3, 2)、#2 の A の形状は (2, 3, 2) であるため、これらがバッチ（0 番目）ディメンションに沿って連結されると、最終的な形状は (3, 3, 2) になります。

tf.matmul が実行され、(3, 3, 4) の形状で出力が生成されます。

tf.matmul の出力はバッチ化されているため、個別のリクエストに分割する必要があります。バッチ処理演算では、各出力テンソルのバッチ（0 番目）ディメンションに沿って分割します。元の入力の形状に基づいて 0 番目のディメンションを分割する方法を決定します。リクエスト #1 の形状は 0 番目のディメンションが 1 であるため、出力は、(1, 3, 4) の形状の 0 番目のディメンション 1 になります。リクエスト #2 の形状は 0 番目のディメンションが 2 であるため、出力は、(2, 3, 4) の形状の 0 番目のディメンション 2 になります。

形状の要件

上記の入力連結と出力分割を実行するために、バッチ処理演算には以下の形状要件があります。

1. バッチ処理への入力をスカラーにすることはできません。0 番目のディメンションに沿って連結するには、テンソルに少なくとも 2 つのディメンションが必要です。

   上のチュートリアルでは、A も B もスカラーではありません。

   この要件を満たさないと、次のようなエラーが発生します: Batching input tensors must have at least one dimension。このエラーの簡単な修正は、スカラーをベクトルにすることです。

2. 異なる推論リクエスト（異なるセッション実行呼び出しなど）間で、同じ名前の入力テンソルは、0 番目のディメンションを除き、各ディメンションに対して同じサイズになります。これにより、入力を 0 番目のディメンションに沿って正常に連結できます。

   上述のチュートリアルでは、リクエスト #1 の A の形状が (1, 3, 2) になっています。 つまり、今後のリクエストでは、(X, 3, 2) のパターンを持つ形状を生成する必要があります。リクエスト #2 は (2, 3, 2) でこの要件を満たしています。同様に、リクエスト #1 の B の形状は (1, 2, 4) であるため、将来のすべてのリクエストで (X, 2, 4) のパターンを持つ形状を生成する必要があります。

   この要件を満たさないと、次のようなエラーが発生します: Dimensions of inputs should match。

3. 特定の推論リクエストでは、すべての入力が同じ 0 番目のディメンション サイズである必要があります。バッチ処理演算への異なる入力テンソルの 0 番目のディメンションが異なる場合、バッチ処理演算では、出力テンソルの分割方法がわかりません。

   上記のチュートリアルでは、リクエスト #1 のテンソルの 0 番目のディメンションサイズはすべて 1 になっています。これにより、バッチ処理では、出力の 0 番目のディメンションサイズが 1 になることがわかります。同様に、リクエスト #2 のテンソルの 0 番目のディメンションサイズは 2 であるため、出力の 0 番目のディメンションサイズは 2 になります。バッチ処演算で (3, 3, 4) の最終形状が分割されると、リクエスト #1 には (1, 3, 4)、リクエスト #2 には (2, 3, 4) が生成されます。

   この要件を満たさないと、次のようなエラーが発生します: Batching input tensors supplied in a given op invocation must have equal 0th-dimension size。

4. 各出力テンソルの形状の 0 番目のディメンションサイズは、すべての入力テンソルの 0 番目のディメンションサイズの合計（および、次に大きい allowed_batch_size を満たすためにバッチ処理オペレーションによって導入されたパディングの合計）でなければなりません。これにより、バッチ処理演算では、入力テンソルの 0 番目のディメンションに基づいて、出力テンソルが 0 番目のディメンションに沿って分割されます。

   上記のチュートリアルでは、入力テンソルにはリクエスト #1 の 0 番目のディメンション 1 と、リクエスト #2 の 0 番目のディメンション 2 があります。したがって、1+2=3 であるため、各出力テンソルの 0 番目のディメンションは 3 である必要があります。出力テンソル (3, 3, 4) は、この要件を満たしています。3 が有効なバッチサイズではなく 4 であった場合、バッチ処理演算では、入力の 0 番目を 3 から 4 までパディングする必要があります。この場合、各出力テンソルの 0 番目のディメンションサイズは 4 である必要があります。

   この要件を満たさないと、次のようなエラーが発生します: Batched output tensor's 0th dimension does not equal the sum of the 0th dimension sizes of the input tensors。

形状要件エラーの解決

これらの要件を満たすには、バッチに別の関数または署名を提供することを検討してください。これらの要件を満たすために既存の関数を変更する必要がある場合もあります。

関数がバッチ化されている場合は、その @tf.function の input_signature の形状がすべて 0 番目のディメンション（バッチ ディメンション）で None になっていることを確認します。署名がバッチ化されている場合は、その入力がすべて、0 番目のディメンションで -1 になっていることを確認します。

BatchFunction 演算は、SparseTensors を入力または出力としてサポートしません。 内部では、それぞれのスパース テンソルは、異なる 0 番目のディメンションサイズを持つことができる 3 つの別個のテンソルとして表されます。