Cloud TPU(TF 2.x)での NCF のトレーニング


概要

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

NCF は、ニューラル ネットワーク アーキテクチャを使用してユーザー アイテムのインタラクションをモデル化するための推奨事項のコラボレーション フィルタリングの一般的なフレームワークです。NCF は、内積をデータから任意の関数を学習できる Multi-Layer Perceptron に置き換えます。

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

以下の手順では、Cloud TPU でモデルをトレーニングすることに精通していることを前提としています。Cloud TPU を初めて使用する場合は、スタートガイドで基本事項をご確認ください。

データセット

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 ユーザーは無料トライアルをご利用いただける場合があります。

始める前に

このチュートリアルを開始する前に、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.

    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.

    Go to project selector

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

  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 を作成するプロジェクトを使用するように Google Cloud CLI を構成します。

    gcloud config set project ${PROJECT_ID}

    このコマンドを新しい Cloud Shell VM で初めて実行すると、[Cloud Shell の承認] ページが表示されます。ページの下部にある [承認] をクリックして、gcloud に認証情報を使用した 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 バケットを作成します。

    gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=europe-west4

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

    バケットのロケーションは、TPU VM と同じリージョンにあることが必要です。TPU VM は、リージョン内の下位区分である特定のゾーンに配置されます。

  6. Cloud TPU VM を作成します。

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

    コマンドフラグの説明

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

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

  7. SSH を使用して Compute Engine インスタンスに接続します。VM に接続すると、シェル プロンプトが username@projectname から username@vm-name に変わります。

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

データを準備する

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

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
  2. データ ディレクトリの環境変数を追加します。

    (vm)$ export DATA_DIR=${STORAGE_BUCKET}/ncf_data
  3. モデルのロケーションを設定し、PYTHONPATH 環境変数を設定します。

    (vm)$ git clone https://github.com/tensorflow/models.git
    (vm)$ pip3 install -r models/official/requirements.txt
    (vm)$ export PYTHONPATH="${PWD}/models:${PYTHONPATH}"
  4. モデル処理ファイルが保存されているディレクトリに移動します。

      (vm)$ cd ~/models/official/recommendation
  5. DATA_DIR で ml-20m データセットのトレーニング データと評価データを生成します。

    (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 名の変数を設定します。

      (vm)$ export TPU_NAME=local

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

次のスクリプトでは、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 リソースを削除します。

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

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

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

次のステップ

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 推論の概要をご覧ください。