Cloud TPU で使用する画像分類データセットの変換

このチュートリアルでは、画像分類データ変換プログラムのサンプル スクリプトを使用して、元の画像分類データセットを Cloud TPU モデルのトレーニングで使用される TFRecord 形式に変換する方法について説明します。

TFRecord では、各画像を個別のファイルとして読み取るよりも Cloud Storage から大きなファイルを読み取るほうが効率的です。TFRecord は、tf.data.Dataset パイプラインを使用している場所であればどこでも使用できます。

TFRecord の使用について詳しくは、次の TensorFlow ドキュメントをご覧ください。

PyTorch または JAX フレームワークを使用しており、データセット ストレージに Cloud Storage を使用していない場合は、TFRecord の利点を活用できない可能性があります。

変換の概要

GitHub のデータ コンバータ リポジトリ内の画像分類フォルダには、converter スクリプト(image_classification_data.py)とサンプル実装(simple_example.py)があります。これらをコピーして変更することで、独自のデータ変換を行うことができます。

画像分類データ コンバータのサンプルでは、ImageClassificationConfigImageClassificationBuilder の 2 つのクラスを定義しています。これらのクラスは tpu/tools/data_converter/image_classification_data.py で定義されています。

ImageClassificationConfig は要約ベースクラスです。ImageClassificationConfig をサブクラス化して、ImageClassificationBuilder をインスタンス化するために必要な構成を定義します。

ImageClassificationBuilder は、画像分類データセットの TensorFlow データセット ビルダーです。これは tdfs.core.GeneratorBasedBuilder のサブクラスです。データセットからデータのサンプルを取得し、TFRecords に変換します。TFRecord は、ImageClassificationBuilder__init__ メソッドの data_dir パラメータで指定されたパスに書き込まれます。

simple_example.py で、サポートされるモードを定義するプロパティ、イメージクラスの数、データセット内の各サンプルの画像データと画像クラスを含む辞書を生成するサンプル生成ツールを実装して、SimpleDatasetConfigImageClassificationConfig をサブクラス化します。

main() 関数は、ランダムに生成された画像データのデータセットを作成し、クラスの数とディスク上のデータセットへのパスを指定する SimpleDatasetConfig オブジェクトをインスタンス化します。次に、main() は、SimpleDatasetConfig インスタンスを渡し、ImageClassificationBuilder オブジェクトをインスタンス化します。最後に、main()download_and_prepare() を呼び出します。このメソッドが呼び出されると、ImageClassificationBuilder インスタンスは SimpleDatasetConfig によって実装されたデータサンプル生成ツールで各サンプルを読み込み、一連の TFRecord ファイルに保存します。

詳細については、分類コンバータのノートブックをご覧ください。

データ コンバージョン サンプルの変更とデータセットの読み込み

データセットを TFRecord 形式に変換するには、次のプロパティを定義する ImageClassificationConfig クラスをサブクラス化します。

  • num_labels: 画像クラスの数を返します
  • supported_modes: データセットでサポートされているモードのリストを返します(例: test、train、validate)
  • text_label_map: テキストクラス ラベルと整数クラスラベル間のマッピングをモデル化する辞書を返します(SimpleDatasetConfig はマッピングを必要としないため、このプロパティを使用しません)。
  • download_path: データセットのダウンロード元のパス(SimpleDatasetConfig はこのプロパティを使用しません。example_generator はディスクからデータを読み込みます)。

example_generator 生成ツール関数を実装します。このメソッドでは、画像データと各サンプルの画像クラス名を含む辞書を生成する必要があります。ImageClassificationBuilderexample_generator() 関数を使用して各サンプルを取得し、TFRecord 形式でディスクに書き込みます。

データ変換サンプルを実行する

  1. 次のコマンドを使用して Cloud Storage バケットを作成します。

    gcloud storage buckets create gs://bucket-name --project=your-project-id --location=us-central2
  2. gcloud コマンドを使用して Cloud TPU を起動します。

    $ gcloud compute tpus tpu-vm create tpu-name \
        --zone=us-central2-b \
        --accelerator-type=v4-8 \
        --version=tpu-vm-tf-2.17.1-pjrt

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン
    accelerator-type
    アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU の各バージョンでサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
    version
    Cloud TPU ソフトウェアのバージョン
  3. SSH を使用して TPU に接続します。

    $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central2-b

    TPU に接続すると、シェル プロンプトが username@projectname から username@vm-name に変わります。

  4. 必要なパッケージをインストールします。

    (vm)$ pip3 install opencv-python-headless pillow
    (vm)$ pip3 install tensorflow-datasets
  5. スクリプトで使用する次の環境変数を作成します。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    (vm)$ export CONVERTED_DIR=$HOME/tfrecords
    (vm)$ export GENERATED_DATA=$HOME/data
    (vm)$ export GCS_CONVERTED=$STORAGE_BUCKET/data_converter/image_classification/tfrecords
    (vm)$ export GCS_RAW=$STORAGE_BUCKET/image_classification/raw
    (vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/tpu/models"
  6. TensorFlow TPU リポジトリをダウンロードします。

    (vm)$ cd /usr/share/
    (vm)$ git clone https://github.com/tensorflow/tpu.git
    (vm)$ cd tpu/tools/data_converter

フェイク データセットでデータ変換ツールを実行する

simple_example.py スクリプトは、データ コンバータ サンプルの image_classification フォルダにあります。次のパラメータを使用してスクリプトを実行すると、一連のフェイク画像が生成され、TFRecord に変換されます。

(vm)$ python3 image_classification/simple_example.py \
  --num_classes=1000 \
  --data_path=$GENERATED_DATA \
  --generate=True \
  --num_examples_per_class_low=10 \
  --num_examples_per_class_high=11 \
  --save_dir=$CONVERTED_DIR

元のデータセットでデータ変換ツールを実行する

  1. 元データの場所の環境変数を作成します。

    (vm)$ export GCS_RAW=gs://cloud-tpu-test-datasets/data_converter/raw_image_classification
  2. simple_example.py スクリプトを実行します。

    (vm)$ python3 image_classification/simple_example.py \
    --num_classes=1000 \
    --data_path=$GCS_RAW \
    --generate=False \
    --save_dir=$CONVERTED_DIR

simple_example.py スクリプトは、次のパラメータを受け取ります。

  • num_classes は、データセット内のクラス数です。ここでは、ImageNet フォーマットに合わせて 1000 を使用しています。
  • generate は、元データを生成するかどうかを決定します。
  • generate=True の場合、data_path はデータが生成されるパス、generate=False の場合は元データが保存されるパスを指します。
  • num_examples_per_class_lownum_examples_per_class_high は、クラスごとにいくつのサンプルを生成するかを決定します。スクリプトは、この範囲内のサンプルの乱数を生成します。
  • save_dir は、保存された TFRecord の保存場所です。Cloud TPU でモデルをトレーニングするには、データは Cloud Storage に保存されている必要があります。データは、Cloud Storage 上でも VM 上でもかまいません。

TFRecord の名前を変更し Cloud Storage に移動する

次の例では、変換されたデータを ResNet モデルで使用します。

  1. TFRecord の名前を ImageNet の TFRecord と同じ形式に変更します。

    (vm)$ cd $CONVERTED_DIR/image_classification_builder/Simple/0.1.0/
    (vm)$ sudo apt update
    (vm)$ sudo apt install rename
    (vm)$ rename -v 's/image_classification_builder-(\w+)\.tfrecord/$1/g' *
  2. TFRecord を Cloud Storage にコピーします。

    (vm)$ gcloud storage cp train* $GCS_CONVERTED
    (vm)$ gcloud storage cp validation* $GCS_CONVERTED

クリーンアップ

  1. Cloud TPU との接続を切断していない場合は切断します。

    (vm)$ exit

    プロンプトが user@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

  2. Cloud Shell で gcloud を実行して、VM リソースを削除します。

    $ gcloud compute tpus tpu-vm delete tpu-name \
      --zone=us-central2-b
  3. gcloud compute tpus tpu-vm list を実行して、VM が削除されたことを確認します。削除には数分かかることがあります。インスタンスが正常に削除されると、次のようなレスポンスが返されます。

    $ gcloud compute tpus tpu-vm list --zone=us-central2-b
    Listed 0 items.
    
  4. 次に示すように gcloud CLI を実行します。bucket-name は、このチュートリアルで作成した Cloud Storage バケット名に置き換えます。

    $ gcloud storage rm gs://bucket-name --recursive