Cloud TPU（TF 2.x）での NCF のトレーニング

---

概要

これは、Neural Matrix Factorization（NeuMF）モデルを使用したニューラル コラボレーション フィルタリング（NCF）フレームワークの実装です。詳細は、ニューラル コラボレーション フィルタリングのレポートで説明されています。現在の実装は、作成者の NCF コードと、MLPerf Repo の Stanford 実装のコードに基づきます。

NCF は、ニューラル ネットワーク アーキテクチャを使用してユーザー インタラクションのインタラクションをモデル化するための推奨事項のコラボレーション フィルタリングの一般的なフレームワークです。従来のモデルとは異なり、NCF はユーザーやアイテムの潜在機能に関する内積がある行列分解（MF）に依存しません。内積は、データから任意の関数を学習できる Multi-Layer Perceptron に置き換えられます。

NCF の 2 つの実装は、Generalized Matrix Factorization（GMF）と Multi-Layer Perceptron（MLP）です。GMF では、潜在機能のインタラクションをモデル化するために線形カーネルを適用し、MLP では非線形カーネルを使用してデータからインタラクション関数を学習します。NFMF は複雑なユーザー - アイテム インタラクションのモデルを改善するための GMF と MLP の融合モデルです。また、ユーザー - アイテムの潜在構造をモデル化するために、MF の線形性と MLF の非線形性の長所を統一します。NFMF を使用すると、GMF と MLP で別々の埋め込みを学習し、最後の隠しレイヤを連結して 2 つのモデルを組み合わせます。neuf_model.py はアーキテクチャの詳細を定義します。

以下の手順では、Cloud TPU でモデルをトレーニングする方法をすでに理解していることを前提としています。Cloud TPU を初めて使用する場合は、クイックスタートで基本的な概要をご確認ください。

警告: NCF では、サードパーティのデータセットを使用します。Google は、このデータセットの有効性またはその他の側面について、いかなる表明、保証もいたしません。

データセット

MovieLens データセットは、モデルのトレーニングと評価に使用されます。具体的には、ml-1m（MovieLens 100 万の略）、ml-20m（MovieLens 2000 万の略）の 2 つのデータセットを使用します。

ml-1m

ml-1m データセットには、2000 年の MovieLens に参加した 6,040 人のユーザーが約 3,706 本の映画を匿名で評価した情報が 1,000,209 件含まれています。すべての評価は、ヘッダー行のない次の形式の ratings.dat ファイルに含まれています。

UserID::MovieID::Rating::Timestamp

• UserID の範囲は 1～6040 です。
• MovieID の範囲は 1～3952 です。
• 評価は 5 段階の星評価です（整数の星評価のみ）。

ml-20m

ml-20m データセットには、138,493 人のユーザーが 26,744 本の映画を評価した情報が 20,000,263 件含まれています。すべての評価が ratings.csv ファイルに含まれています。このファイルでは、ヘッダー行に続く各行に、単一のユーザーの 1 本の映画に対する評価が次の形式で 1 件ずつ記録されています。

userId,movieId,rating,timestamp

このファイル内の行は、まず userId の順に並べられ、次にユーザー内で movieId の順に並べられます。評価は 5 段階の星評価で、0.5 星ごとに増分します（0.5～5.0 星）。どちらのデータセットも、協定世界時（UTC）の1970 年 1 月 1 日午前 0 時を起点としたタイムスタンプが秒単位で表されます。各ユーザーには 20 以上の評価があります。

目標

• データセットとモデルの出力を格納する Cloud Storage バケットを作成する
• MovieLens データセットを準備する
• トレーニングと評価のための Compute Engine VM と Cloud TPU ノードを設定する
• トレーニングと評価を行う

費用

このドキュメントでは、Google Cloud の次の課金対象のコンポーネントを使用します。

• Compute Engine
• Cloud TPU
• Cloud Storage

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

始める前に

重要: このチュートリアルは、TPU VM または TPU ノード構成で使用できます。この 2 つの VM アーキテクチャについては、システム アーキテクチャをご覧ください。使用する gcloud コマンドは、使用する TPU 構成によって異なります。このチュートリアルでは、それぞれの gcloud コマンドをタブ付きセクションに示します。使用する TPU 構成のタブを選択すると、対応する gcloud コマンドがウェブページに表示されます。TPU ノードを使用する必要がある場合を除き、TPU VM を使用することをおすすめします。

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

1. 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.

2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

3. Make sure that billing is enabled for your Google Cloud project.

4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

5. Make sure that billing is enabled for your Google Cloud project.

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

リソースを設定する

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

重要: ネットワークのレイテンシを抑え費用を節減するために、すべてのリソースを同じリージョン / ゾーンに設定してください。

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

   Cloud Shell を開く

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

   export PROJECT_ID=project-id

3. Cloud TPU を作成するプロジェクトを使用するように Google Cloud CLI を構成します。

   gcloud config set project ${PROJECT_ID}

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

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

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

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

   [e-mail]
   

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

   注: 次のコマンドの bucket-name は、バケットに割り当てる名前に置き換えてください。

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

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

   バケットのロケーションは、仮想マシン（VM）および TPU ノードと同じリージョンにする必要があります。VM と TPU ノードは、リージョン内のサブディビジョンである特定のゾーンに配置されます。

6. gcloud コマンドを使用して Compute Engine VM と Cloud TPU を起動します。使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。2 つの VM アーキテクチャの詳細については、システム アーキテクチャをご覧ください。

   TPU VM

   $ gcloud compute tpus tpu-vm create ncf-tutorial \
   --zone=europe-west4-a \
   --accelerator-type=v3-8 \
   --version=tpu-vm-tf-2.16.1-pjrt

   コマンドフラグの説明

   zone

   Cloud TPU を作成するゾーン。

   accelerator-type

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

   version

   Cloud TPU ソフトウェアのバージョン。

   TPU ノード

   $ gcloud compute tpus execution-groups create  \
    --zone=europe-west4-a \
    --name=ncf-tutorial \
    --accelerator-type=v3-8 \
    --machine-type=n1-standard-8 \
    --disk-size=300 \
    --tf-version=2.12.0

   コマンドフラグの説明

   zone

   Cloud TPU を作成するゾーン。

   name

   TPU 名。指定しない場合、デフォルトでユーザー名が使用されます。

   accelerator-type

   作成する Cloud TPU のタイプ。

   machine-type

   作成する Compute Engine VM のマシンタイプ。

   disk-size

   Compute Engine VM のルート ボリューム サイズ（GB）。

   tf-version

   Tensorflow gcloud のバージョンが VM にインストールされます。

   注: 複数のプロジェクトがある場合は、--project フラグでプロジェクト ID を指定する必要があります。

   gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。

   注: プロジェクトで初めて gcloud compute tpus を実行する場合、SSH 認証鍵の伝播や API の始動などの起動タスクの実行に約 5 分かかります。

7. 自動的に Compute Engine インスタンスにログインしない場合は、次の ssh コマンドを実行してログインします。VM にログインすると、シェル プロンプトが username@projectname から username@vm-name に変わります。

   TPU VM

   gcloud compute tpus tpu-vm ssh ncf-tutorial --zone=europe-west4-a

   TPU ノード

   gcloud compute ssh ncf-tutorial --zone=europe-west4-a

   重要なポイント: これ以降、接頭辞 (vm) $ は Compute Engine VM インスタンスでコマンドを実行する必要があることを意味します。

データを準備する

1. ストレージ バケットの環境変数を追加します。bucket-name は、使用するバケット名に置き換えてください。

   (vm)$ export STORAGE_BUCKET=gs://bucket-name

2. データ ディレクトリの環境変数を追加します。

   (vm)$ export DATA_DIR=${STORAGE_BUCKET}/ncf_data

3. モデルのロケーションを設定し、PYTHONPATH 環境変数を設定します。

   TPU VM

   (vm)$ git clone https://github.com/tensorflow/models.git
   (vm)$ pip3 install -r models/official/requirements.txt

   (vm)$ export PYTHONPATH="${PWD}/models:${PYTHONPATH}"

   TPU ノード

   (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
   (vm)$ pip3 install -r /usr/share/models/official/requirements.txt

4. モデル処理ファイルが保存されているディレクトリに移動します。

   TPU VM

   (vm)$ cd ~/models/official/recommendation

   TPU ノード

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

5. DATA_DIR で ml-20m データセットのトレーニング データと評価データを生成します。

   注: Cloud Shell の接続が切断されたときにダウンロードが終了することを回避するために、ダウンロード コマンドの前に nohup を付けるか、screen を使用することをおすすめします。

   (vm)$ python3 create_ncf_data.py \
       --dataset ml-20m \
       --num_train_epochs 4 \
       --meta_data_file_path ${DATA_DIR}/metadata \
       --eval_prebatch_size 160000 \
       --data_dir ${DATA_DIR}

このスクリプトは、VM 上でデータセットを生成して前処理します。前処理を行うと、データがモデルに必要な TFRecord 形式に変換されます。ダウンロードと前処理には約 25 分かかり、次のような出力が生成されます。

I0804 23:03:02.370002 139664166737728 movielens.py:124] Successfully downloaded /tmp/tmpicajrlfc/ml-20m.zip 198702078 bytes
I0804 23:04:42.665195 139664166737728 data_preprocessing.py:223] Beginning data preprocessing.
I0804 23:04:59.084554 139664166737728 data_preprocessing.py:84] Generating user_map and item_map...
I0804 23:05:20.934210 139664166737728 data_preprocessing.py:103] Sorting by user, timestamp...
I0804 23:06:39.859857 139664166737728 data_preprocessing.py:194] Writing raw data cache.
I0804 23:06:42.375952 139664166737728 data_preprocessing.py:262] Data preprocessing complete. Time: 119.7 sec.
%lt;BisectionDataConstructor(Thread-1, initial daemon)>
General:
  Num users: 138493
  Num items: 26744

Training:
  Positive count:          19861770
  Batch size:              99000
  Batch count per epoch:   1004

Eval:
  Positive count:          138493
  Batch size:              160000
  Batch count per epoch:   866

I0804 23:07:14.137242 139664166737728 data_pipeline.py:887] Negative total vector built. Time: 31.8 seconds
I0804 23:11:25.013135 139664166737728 data_pipeline.py:588] Epoch construction complete. Time: 250.9 seconds
I0804 23:15:46.391308 139664166737728 data_pipeline.py:674] Eval construction complete. Time: 261.4 seconds
I0804 23:19:54.345858 139664166737728 data_pipeline.py:588] Epoch construction complete. Time: 248.0 seconds
I0804 23:24:09.182484 139664166737728 data_pipeline.py:588] Epoch construction complete. Time: 254.8 seconds
I0804 23:28:26.224653 139664166737728 data_pipeline.py:588] Epoch construction complete. Time: 257.0 seconds

Cloud TPU を設定して起動する

1. Cloud TPU 名の変数を設定します。

   TPU VM

   (vm)$ export TPU_NAME=local

   TPU ノード

   (vm)$ export TPU_NAME=ncf-tutorial

トレーニングと評価を行う

次のスクリプトでは、3 つのエポック用のサンプル トレーニングを実行します。

1. チェックポイントと TensorBoard のサマリーを保存するための環境変数をモデルのディレクトリに追加します。

   (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/ncf

2. TPU を作成するときに、--version パラメータを -pjrt で終わるバージョンに設定した場合は、次の環境変数を設定して PJRT ランタイムを有効にします。

     (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
     (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

3. 次のコマンドを実行して、NCF モデルをトレーニングします。

   (vm)$ python3 ncf_keras_main.py \
        --model_dir=${MODEL_DIR} \
        --data_dir=${DATA_DIR} \
        --train_dataset_path=${DATA_DIR}/training_cycle_*/* \
        --eval_dataset_path=${DATA_DIR}/eval_data/* \
        --input_meta_data_path=${DATA_DIR}/metadata \
        --learning_rate=3e-5 \
        --train_epochs=3 \
        --dataset=ml-20m \
        --eval_batch_size=160000 \
        --learning_rate=0.00382059 \
        --beta1=0.783529 \
        --beta2=0.909003 \
        --epsilon=1.45439e-07 \
        --dataset=ml-20m \
        --num_factors=64 \
        --hr_threshold=0.635 \
        --keras_use_ctl=true \
        --layers=256,256,128,64 \
        --use_synthetic_data=false \
        --distribution_strategy=tpu \
        --download_if_missing=false
    

トレーニングと評価には約 2 分かかり、次のような最終出力が生成されます。

Result is {'loss': <tf.Tensor: shape=(), dtype=float32, numpy=0.10950611>,
'train_finish_time': 1618016422.1377568, 'avg_exp_per_second': 3062557.5070816963}

クリーンアップ

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

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

   (vm)$ exit

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

2. Cloud TPU と Compute Engine リソースを削除します。リソースの削除に使用するコマンドは、TPU VM または TPU ノードのどちらを使用するかによって異なります。詳細については、システム アーキテクチャをご覧ください。

   TPU VM

   $ gcloud compute tpus tpu-vm delete ncf-tutorial \
   --zone=europe-west4-a

   TPU ノード

   $ gcloud compute tpus execution-groups delete ncf-tutorial \
   --zone=europe-west4-a

3. 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

   Listed 0 items.
   

4. 次に示すように gsutil を実行します。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 推論の概要をご覧ください。

• ローカルマシンへのインストール方法など、ctpu の詳細を学習する。
• TensorBoard の TPU ツールを確認します。