Cloud TPU(TF 2.x)での DLRM および DCN のトレーニング

このチュートリアルでは、クリック率(CTR)予測などのタスクに使用できる DLRM と DCN v2 のランキング モデルをトレーニングする方法について説明します。DLRM または DCN v2 ランキング モデルをトレーニングするパラメータの設定方法については、DLRM または DCN モデルを実行するための設定をご覧ください。

モデル入力は数値特徴とカテゴリ特徴であり、出力はスカラー(クリック確率など)です。モデルは、Cloud TPU でトレーニングおよび評価できます。ディープ ランキング モデルは、メモリ消費量が多く(テーブル / ルックアップの埋め込みのため)、ディープ ネットワーク(MLP)のためにコンピューティング負荷が高くなります。TPU は両方の用途向けに設計されています。

このモデルでは、カテゴリ特徴に TPUEmbedding レイヤを使用します。TPU 埋め込みは、高速ルックアップを使用して大きな埋め込みテーブルをサポートします。埋め込みテーブルのサイズは、TPU Pod のサイズに比例して増減します。TPU v3-8 では最大 90 GB の埋め込みテーブルを使用できます。v3-512 Pod では最大 5.6 TB、v3-2048 TPU Pod では最大 22.4 TB を使用できます。

モデルコードは TensorFlow Recommenders ライブラリにありますが、入力パイプライン、構成、トレーニング ループについては TensorFlow Model Garden をご覧ください。

目標

  • トレーニング環境を設定する
  • 合成データを使用してトレーニング ジョブを実行する
  • 出力結果を確認します。

費用

このチュートリアルでは、課金対象である次の Google Cloud コンポーネントを使用します。

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。 新しい Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

始める前に

このチュートリアルを開始する前に、Google Cloud プロジェクトが正しく設定されていることを確認します。

  1. Google Cloud アカウントにログインします。Google Cloud を初めて使用する場合は、アカウントを作成して、実際のシナリオでの Google プロダクトのパフォーマンスを評価してください。新規のお客様には、ワークロードの実行、テスト、デプロイができる無料クレジット $300 分を差し上げます。
  2. Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

    プロジェクト セレクタに移動

  3. Cloud プロジェクトに対して課金が有効になっていることを確認します。プロジェクトに対して課金が有効になっていることを確認する方法を学習する

  4. Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

    プロジェクト セレクタに移動

  5. Cloud プロジェクトに対して課金が有効になっていることを確認します。プロジェクトに対して課金が有効になっていることを確認する方法を学習する

  6. このチュートリアルでは、Google Cloud の課金対象となるコンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。不要な課金を回避するために、このチュートリアルを完了したら、作成したリソースを必ずクリーンアップしてください。

リソースを設定する

このセクションでは、このチュートリアルで使用する Cloud Storage のバケット、VM、Cloud TPU の各リソースを設定する方法について説明します。

  1. Cloud Shell ウィンドウを開きます。

    Cloud Shell を開く

  2. プロジェクト ID の変数を作成します。

    export PROJECT_ID=project-id
    
  3. Cloud TPU を作成するプロジェクトを使用するように gcloud コマンドライン ツールを構成します。

    gcloud config set project ${PROJECT_ID}
    

    このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した GCP API の呼び出しを許可します。

  4. Cloud TPU プロジェクトのサービス アカウントを作成します。

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
    

    このコマンドでは、Cloud TPU サービス アカウントを次の形式で返します。

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

  5. 次のコマンドを使用して Cloud Storage バケットを作成します。ここで、-l オプションは、バケットを作成するリージョンを指定します。ゾーンとリージョンの詳細については、タイプとゾーンをご覧ください。

    gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 gs://bucket-name
    

    この Cloud Storage バケットには、モデルのトレーニングに使用するデータとトレーニング結果が格納されます。このチュートリアルで使用する gcloud compute tpus execution-groups ツールは、前の手順で設定した Cloud TPU サービス アカウントのデフォルトの権限を設定します。権限の詳細な設定が必要な場合は、アクセスレベル権限をご覧ください。

    バケットのロケーションは、Compute Engine(VM)および Cloud TPU ノードと同じリージョンにする必要があります。

  6. gcloud コマンドを使用して Compute Engine VM と Cloud TPU を起動します。

    $ gcloud compute tpus execution-groups create \
     --vm-only \
     --name=dlrm-dcn-tutorial \
     --zone=zone \
     --disk-size=300 \
     --machine-type=n1-standard-8 \
     --tf-version=2.7.0
    

    コマンドフラグの説明

    vm-only
    VM のみを作成します。デフォルトでは、gcloud compute tpus execution-groups コマンドは VM と Cloud TPU を作成します。
    name
    作成する Cloud TPU の名前。
    zone
    Cloud TPU を作成するゾーン
    disk-size
    gcloud compute tpus execution-groups コマンドで作成された VM のハードディスクのサイズ(GB)。
    machine-type
    作成する Compute Engine VM のマシンタイプ
    tf-version
    Tensorflow ctpu のバージョンが VM にインストールされます。

    gcloud compute tpus execution-groups コマンドの実行が終了したら、shell プロンプトが username@projectname から username@vm-name に変更されたことを確認します。変更されていれば、Compute Engine VM にログインしていることになります。

    Compute Engine インスタンスに接続していない場合は、次のコマンドを実行して接続できます。

    gcloud compute ssh dlrm-dcn-tutorial --zone=zone
    

    これらの手順を続行する場合は、Compute Engine インスタンスで (vm)$ で始まる各コマンドを実行します。

Cloud Storage バケットの変数を設定する

次の環境変数を設定します。bucket-name を Cloud Storage バケットの名前に置き換えます。

(vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export TPU_NAME=dlrm-dcn-tutorial
(vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/models/"
(vm)$ export EXPERIMENT_NAME=dlrm-exp

トレーニング アプリケーションでは、Cloud Storage でトレーニング データにアクセスできる必要があります。また、トレーニング アプリケーションは、Cloud Storage バケットを使用してトレーニング中にチェックポイントを保存します。

合成データで DLRM または DCN モデルを実行するように設定する

モデルは、さまざまなデータセットでトレーニングできます。Criteo TerabyteCriteo Kaggle の 2 つが一般的に使用されます。このチュートリアルでは、use_synthetic_data=True フラグを設定して合成データをトレーニングします。

合成データセットは、Cloud TPU の使用方法を理解し、エンドツーエンドのパフォーマンスを検証する場合にのみ役立ちます。精度の数値と保存されたモデルは意味がありません。

これらのデータセットをダウンロードして前処理する方法については、Criteo Terabyte および Criteo Kaggle ウェブサイトにアクセスしてください。

  1. gcloud コマンドを使用して Cloud TPU リソースを起動します。

    (vm)$ gcloud compute tpus execution-groups create \
     --tpu-only \
     --accelerator-type=v3-8  \
     --name=dlrm-dcn-tutorial \
     --zone=zone \
     --tf-version=2.7.0
    

    コマンドフラグの説明

    tpu-only
    VM を作成せずに Cloud TPU を作成します。デフォルトでは、gcloud compute tpus execution-groups コマンドは VM と Cloud TPU を作成します。
    accelerator-type
    作成する Cloud TPU のタイプ
    name
    作成する Cloud TPU の名前。
    zone
    Cloud TPU を作成するゾーン
    tf-version
    Tensorflow gcloud のバージョンが VM にインストールされます。
  2. 必要なパッケージをインストールします。

    (vm)$ pip install tensorflow-recommenders
    
  3. トレーニング スクリプトを実行します。これは、架空の Criteo に似たデータセットを使用して、DLRM モデルをトレーニングします。トレーニングには約 20 分かかります。

python3 /usr/share/models/official/recommendation/ranking/train.py --mode=train_and_eval \
--model_dir=${STORAGE_BUCKET}/model_dirs/${EXPERIMENT_NAME} --params_override="
runtime:
    distribution_strategy: 'tpu'
task:
    use_synthetic_data: true
    train_data:
        input_path: '${DATA_DIR}/train/*'
        global_batch_size: 16384
    validation_data:
        input_path: '${DATA_DIR}/eval/*'
        global_batch_size: 16384
    model:
        num_dense_features: 13
        bottom_mlp: [512,256,64]
        embedding_dim: 64
        top_mlp: [1024,1024,512,256,1]
        interaction: 'dot'
        vocab_sizes: [39884406, 39043, 17289, 7420, 20263, 3, 7120, 1543, 63,
            38532951, 2953546, 403346, 10, 2208, 11938, 155, 4, 976, 14,
            39979771, 25641295, 39664984, 585935, 12972, 108, 36]
trainer:
    use_orbit: false
    validation_interval: 1000
    checkpoint_interval: 1000
    validation_steps: 500
    train_steps: 1000
    steps_per_loop: 1000
"

コマンドフラグの説明

use_orbit
トレーニングに Orbit ライブラリを使用するか、keras のコンパイル/フィット API を使用するかを指定します。
train_steps
モデルをトレーニングするためのステップの数。
validation_steps
評価の実行に使用されるステップの数。
validation_interval
評価間で実行するトレーニング ステップの数は、train_steps 以下にする必要があります。
steps_per_loop
グラフモード ループごとのステップ数。これにより、eager コンテキストでの通信が減少します。

このトレーニングは、v3-8 TPU で約 10 分間実行されます。完了すると、次のようなメッセージが表示されます。

I0621 21:32:58.519792 139675269142336 tpu_embedding_v2_utils.py:907] Done with log of TPUEmbeddingConfiguration.
I0621 21:32:58.540874 139675269142336 tpu_embedding_v2.py:389] Done initializing TPU Embedding engine.
1000/1000 [==============================] - 335s 335ms/step - auc: 0.7360 - accuracy: 0.6709 - prediction_mean: 0.4984
- label_mean: 0.4976 - loss: 0.0734 - regularization_loss: 0.0000e+00 - total_loss: 0.0734 - val_auc: 0.7403
- val_accuracy: 0.6745 - val_prediction_mean: 0.5065 - val_label_mean: 0.4976 - val_loss: 0.0749
- val_regularization_loss: 0.0000e+00 - val_total_loss: 0.0749

Model: "ranking"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
tpu_embedding (TPUEmbedding) multiple                  1
_________________________________________________________________
mlp (MLP)                    multiple                  154944
_________________________________________________________________
mlp_1 (MLP)                  multiple                  2131969
_________________________________________________________________
dot_interaction (DotInteract multiple                  0
_________________________________________________________________
ranking_1 (Ranking)          multiple                  0
=================================================================
Total params: 2,286,914
Trainable params: 2,286,914
Non-trainable params: 0
_________________________________________________________________
I0621 21:43:54.977140 139675269142336 train.py:177] Train history: {'auc': [0.7359596490859985],
'accuracy': [0.67094486951828], 'prediction_mean': [0.4983849823474884], 'label_mean': [0.4975697994232178],
'loss': [0.07338511198759079], 'regularization_loss': [0], 'total_loss': [0.07338511198759079],
'val_auc': [0.7402724623680115], 'val_accuracy': [0.6744520664215088], 'val_prediction_mean': [0.5064718723297119],
'val_label_mean': [0.4975748658180237], 'val_loss': [0.07486172765493393],
'val_regularization_loss': [0], 'val_total_loss': [0.07486172765493393]}

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

  1. Compute Engine インスタンスとの接続を切断していない場合は切断します。

    (vm)$ exit
    

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

  2. Cloud Shell で次のコマンドを使用して、Compute Engine VM と Cloud TPU を削除します。

    $ gcloud compute tpus execution-groups delete dlrm-dcn-tutorial \
      --zone=zone
    
  3. gcloud compute tpus execution-groups list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。以下のようなレスポンスは、インスタンスが正常に削除されたことを示します。

    $ gcloud compute tpus execution-groups list \
     --zone=zone
    

    次のような空の TPU のリストが表示されます。

       NAME             STATUS
    
  4. 次に示すように、gsutil を使用して Cloud Storage バケットを削除します。bucket-name を Cloud Storage バケットの名前に置き換えます。

    $ gsutil rm -r gs://bucket-name
    

次のステップ

TensorFlow Cloud TPU のチュートリアルでは通常、サンプル データセットを使用してモデルをトレーニングします。このトレーニングの結果は(ほとんどの場合)推論には使用できません。モデルを推論に使用するには、一般公開されているデータセットまたは独自のデータセットでデータをトレーニングします。Cloud TPU でトレーニングされた TensorFlow モデルは通常、データセットを TFRecord 形式にする必要があります。

データセット変換ツールのサンプルを使用して、画像分類データセットを TFRecord 形式に変換できます。画像分類モデルを使用しない場合は、自分でデータセットを TFRecord 形式に変換する必要があります。詳細については、TFRecord と tf.Example をご覧ください。

ハイパーパラメータ調整

データセットでモデルのパフォーマンスを向上させるには、モデルのハイパーパラメータを調整します。すべての TPU でサポートされているモデルに共通のハイパーパラメータに関する情報については、GitHub をご覧ください。モデルに固有のハイパーパラメータに関する情報については、各モデルのソースコードで確認できます。ハイパーパラメータ調整の詳細については、ハイパーパラメータ調整の概要ハイパーパラメータ調整サービスの使用ハイパーパラメータを調整するをご覧ください。

推論

モデルをトレーニングしたら、そのモデルを推論(予測)に使用できます。AI Platform は、機械学習モデルを開発、トレーニングデプロイするためのクラウドベースのソリューションです。モデルをデプロイすれば、AI Platform Prediction サービスを使用できるようになります。