Cloud TPU による大きな画像の処理

このドキュメントでは、Cloud TPU と TPUEstimator による空間パーティショニングを使用して、非常に大きな画像と動画のトレーニングを可能にする方法を説明します。

Cloud TPU の空間パーティショニング機能を使用すると、1 つの画像を分割して複数の TPU チップに分散できます。空間パーティショニングにより、アクセラレータ チップ上の 1 つのコアで使用可能なメモリに収まりきらない大きな画像でも、入力画像として使用して実行できるよう、簡単にモデルをスケーリングできます。

この機能はアクティベーションを複数の TPU コアに分散するため、トレーニングに 2 個、4 個、8 個、さらには 16 個のコアを使用するようにモデルをスケーリングできます。トレーニングに使用できるプロセッサの数をスケーリングすると、画像データのダウンサンプリングが不要になり、精度が高くなるだけでなく、トレーニング パフォーマンスの効率も向上します。

空間パーティショニングは、次のようなタスクに適用できます。

  • 3D コンピュータ断層撮影（CT スキャン）画像のセグメンテーション
  • 動画のコンテンツ分析
  • 画像データをダウンサンプリングしない、自律駆動型のオブジェクト検出

TPUEstimator による空間パーティショニングの有効化

Tensorflow では、Tensorflow を最適化する XLA コンパイラがすべての Cloud TPU コア間の通信を自動的に処理します。モデルの空間パーティショニングを有効にするためにコードに変更を加える必要はありません。TPUEstimator は空間パーティショニング API をサポートしているため、TPUConfig 内で各入力テンソルのパーティショニング方法を構成するだけで、空間パーティショニングを有効にできます。

画像

例

次のコードは、画像分類モデルの 4 方向空間パーティショニングを行う TPUConfig の例を示しています。テンソルは、高さの次元に沿って 4 つの部分に分割されます（テンソルの形状が [batch, height, width, channel] であると想定した場合）。

tpu_config=tpu_config.TPUConfig(
    iterations_per_loop=100,
    num_cores_per_replica=4,
    per_host_input_for_training=tpu_config.InputPipelineConfig.PER_HOST_V2,
    input_partition_dims=[[1, 4, 1, 1], None]]

空間パーティショニングでは、入力パイプラインの形式が tf.data であること、per_host_input_for_training の値（train_batch_size）が PER_HOST_V2 に設定されていることが必要です。指定する num_cores_per_replica によって、画像を分割できるパーティションの最大数が決まります。

input_partition_dims の値は、feature_partition_dims と label_partition_dims の 2 要素からなるリストを提供し、この 2 つの要素によって入力テンソルのパーティショニング方法が記述されます。feature_partition_dims と label_partition_dims の構造は、input_fn の機能とラベルの構造と一致している必要があります。ラベル パーティションに「None」を設定すると、ラベルは分割されません。

空間パーティションを使用したリファレンス モデルの実行

2D オブジェクト検出

RetinaNet は、画像内のオブジェクトを境界ボックスで特定し、識別したオブジェクトを分類する、オブジェクト検出モデルです。1 つの Cloud TPU コア（デバイスあたり 8 個のバッチ）で処理できる画像の最大サイズは 1280x1280 であるため、非常に大きな画像を使用してこのモデルをトレーニングするのは困難です。モデルの空間パーティショニングを使用して 8 個の TPU コアで分散処理すれば、最大サイズの 4 倍大きい画像でトレーニングできます。

画像サイズ TPU タイプ TPU コア数 グローバル バッチサイズ ステップ時間
1280x1280 v3-8 8 64 910 ミリ秒
2560x2560 v3-64 64 64 822 ミリ秒

3D 画像セグメンテーション

3D UNet は、医用画像分野で広く使用されている高密度 3D セグメンテーション モデルです。CT 画像の元の解像度は 256x256x256 まで高くなることもあり、1 つの TPU コアでは処理しきれません。したがって、通常は画像のダウンサンプリングが必要になります。TPU 空間パーティショニングを使用すれば、16 方向の空間パーティショニングを使って元の解像度のまま画像を処理できます。

画像サイズ TPU タイプ TPU コア数 グローバル バッチサイズ ステップ時間
128x128x128 v3-8 8 32 3.428 秒
2560x2560 v3-64 64 32 3.02 秒

参考情報