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 プロジェクトが正しく設定されていることを確認します。

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

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

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

Google Cloud プロジェクトで課金が有効になっていることを確認します。

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

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

Google Cloud プロジェクトで課金が有効になっていることを確認します。

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

リソースを設定する

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

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

Cloud Shell を開く
プロジェクト ID の変数を作成します。
```
export PROJECT_ID=project-id
```
Cloud TPU を作成するプロジェクトを使用するように Google Cloud CLI を構成します。

注: gcloud コマンドの詳細については、Google Cloud CLI リファレンスをご覧ください。
```
gcloud config set project ${PROJECT_ID}
```
このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した API の呼び出しを許可します。
Cloud TPU プロジェクトのサービスアカウントを作成します。
```
gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
```
このコマンドでは、Cloud TPU サービスアカウントを次の形式で返します。
```
service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
```
次のコマンドを使用して 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 ノードと同じリージョンにする必要があります。
gcloud コマンドを使用して Compute Engine VM と Cloud TPU を起動します。使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。2 つの VM アーキテクチャの詳細については、システムアーキテクチャをご覧ください。gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。
TPU VM
```
$ gcloud compute tpus tpu-vm create dlrm-dcn-tutorial \
--zone=europe-west4-a \
--accelerator-type=v3-8 \
--version=tpu-vm-tf-2.16.1-se
```
コマンドフラグの説明

zone

Cloud TPU を作成するゾーン。

accelerator-type

アクセラレータタイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータタイプの詳細については、TPU のバージョンをご覧ください。

version

Cloud TPU ソフトウェアのバージョン。
TPU ノード
```
$ gcloud compute tpus execution-groups create \
--name=dlrm-dcn-tutorial \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-8 \
--tf-version=2.12.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 を実行する場合、SSH 認証鍵の伝播や API の始動などの起動タスクの実行に数分かかる場合があります。
自動的に Compute Engine インスタンスにログインしない場合は、次の ssh コマンドを実行してログインします。VM にログインすると、シェルプロンプトが username@projectname から username@vm-name に変わります。
TPU VM
```
gcloud compute tpus tpu-vm ssh dlrm-dcn-tutorial --zone=europe-west4-a
```
TPU ノード
```
gcloud compute ssh dlrm-dcn-tutorial --zone=europe-west4-a
```
これらの手順を続行する場合は、VM セッションウィンドウで、(vm)$ で始まる各コマンドを実行します。

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

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

(vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export PYTHONPATH="/usr/share/tpu/models/:${PYTHONPATH}"
(vm)$ export EXPERIMENT_NAME=dlrm-exp

TPU 名の環境変数を設定します。

TPU VM

(vm)$ export TPU_NAME=local

TPU ノード

(vm)$ export TPU_NAME=dlrm-dcn-tutorial

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

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

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

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

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

必要なパッケージをインストールする

(vm)$ pip3 install tensorflow-recommenders
(vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt

スクリプトディレクトリに移動します。

TPU VM

(vm)$ cd /usr/share/tpu/models/official/recommendation/ranking

TPU ノード

(vm)$ cd /usr/share/models/official/recommendation/ranking

トレーニングスクリプトを実行します。これは、架空の Criteo に似たデータセットを使用して、DLRM モデルをトレーニングします。トレーニングには約 20 分かかります。

export EMBEDDING_DIM=32

python3 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,${EMBEDDING_DIM}]
             embedding_dim: ${EMBEDDING_DIM}
             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
     "

このトレーニングは、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 アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

Compute Engine インスタンスとの接続を切断していない場合は切断します。
```
(vm)$ exit
```
プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。
Cloud TPU と Compute Engine リソースを削除します。リソースの削除に使用するコマンドは、TPU VM または TPU ノードのどちらを使用するかによって異なります。詳細については、システムアーキテクチャをご覧ください。
TPU VM
```
$ gcloud compute tpus tpu-vm delete dlrm-dcn-tutorial \
--zone=europe-west4-a
```
TPU ノード
```
$ gcloud compute tpus execution-groups delete dlrm-dcn-tutorial \
--zone=europe-west4-a
```
gcloud compute tpus execution-groups list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。次のコマンドの出力には、このチュートリアルで作成したリソースを含めないでください。
TPU VM
```
$ gcloud compute tpus tpu-vm list --zone=europe-west4-a
```
TPU ノード
```
$ gcloud compute tpus execution-groups list --zone=europe-west4-a
```
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 をご覧ください。モデルに固有のハイパーパラメータに関する情報については、各モデルのソースコードで確認できます。ハイパーパラメータ調整の詳細については、ハイパーパラメータ調整の概要とハイパーパラメータを調整するをご覧ください。

推論

モデルをトレーニングしたら、そのモデルを推論（予測）に使用できます。Cloud TPU 推論コンバータツールを使用して、Cloud TPU v5e での推論用の TensorFlow モデルを準備して最適化できます。Cloud TPU v5e での推論の詳細については、Cloud TPU v5e 推論の概要をご覧ください。

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

目標

費用

始める前に

リソースを設定する

TPU VM

コマンドフラグの説明

TPU ノード

コマンドフラグの説明

TPU VM

TPU ノード

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

TPU VM

TPU ノード

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

TPU VM

TPU ノード

クリーンアップ

TPU VM

TPU ノード

TPU VM

TPU ノード

次のステップ

ハイパーパラメータ チューニング

推論

ハイパーパラメータチューニング