Cloud TPU での PyTorch を使用した Wav2Vec2 の事前トレーニング

このチュートリアルでは、PyTorch を使用して Cloud TPU デバイスで FairSeq の Wav2Vec2 モデルを事前トレーニングする方法を説明します。PyTorch と ImageNet データセットを使用する、TPU 用に最適化されたその他のイメージ分類モデルにも、同じパターンを適用できます。

このチュートリアルのモデルは、wav2vec 2.0: Framework for Self-Supervised Learning of Speech Representations の論文に基づいています。

目標

PyTorch 環境を作成して構成します。
オープンソース LibriSpeech データをダウンロードします。
トレーニングジョブを実行します。

費用

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

Compute Engine
Cloud TPU

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

始める前に

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

Google Cloud Console のプロジェクトセレクタページで、Google Cloud プロジェクトを選択または作成します。注: この手順で作成するリソースを保持し続けない場合は、既存のプロジェクトを選択するのではなく、新しいプロジェクトを作成してください。チュートリアルが終了した後、そのプロジェクトを削除して、プロジェクトに関連するすべてのリソースを削除できます。
プロジェクトを選択するページに移動し、Cloud プロジェクトに対して課金が有効になっていることを確認します。プロジェクトに対して課金が有効になっていることを確認する方法を確かめる。

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

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

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

Cloud プロジェクトに対して課金が有効になっていることを確認します。詳しくは、プロジェクトで課金が有効になっているかどうかを確認する方法をご覧ください。

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

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

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

Compute Engine インスタンスを設定する

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

Cloud Shell を開く
プロジェクト ID の変数を作成します。
```
export PROJECT_ID=project-id
```
Cloud TPU を作成するプロジェクトを使用するように Google Cloud CLI を構成します。
```
gcloud config set project ${PROJECT_ID}
```
このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した API の呼び出しを許可します。
Cloud Shell から、このチュートリアルに必要な Compute Engine リソースを起動します。

注: Criteo Kaggle データセットでトレーニングする場合は、n1-highmem-96 マシンタイプを使用してください。
```
gcloud compute instances create wav2vec2-tutorial \
  --zone=us-central1-a \
  --machine-type=n1-standard-64 \
  --image-family=torch-xla \
  --image-project=ml-images  \
  --boot-disk-size=200GB \
  --scopes=https://www.googleapis.com/auth/cloud-platform
```
注:
このコマンドを実行すると、次の警告が表示されます。
```
Disk size: '200 GB' is larger than image size
```
この警告メッセージは無視してかまいません。
新しい Compute Engine インスタンスに接続します。
```
gcloud compute ssh wav2vec2-tutorial --zone=us-central1-a
```

Cloud TPU リソースを起動する

Compute Engine 仮想マシンで、PyTorch バージョンを設定します。
```
(vm) $ export PYTORCH_VERSION=1.13
```

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

(vm) $ gcloud compute tpus create wav2vec2-tutorial \
--zone=us-central1-a \
--network=default \
--version=pytorch-1.13 \
--accelerator-type=v3-8

次のコマンドを実行して、Cloud TPU リソースの IP アドレスを識別します。
```
(vm) $ gcloud compute tpus describe wav2vec2-tutorial --zone=us-central1-a
```
このコマンドにより、TPU に関する情報が表示されます。ipAddress というラベルの付いた行を探します。次のコマンドで、ip-address を IP アドレスに置き換えます。
```
(vm) $ export TPU_IP_ADDRESS=ip-address
(vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
```

PyTorch 環境を作成および構成する

conda 環境を開始します。
```
(vm) $ conda activate torch-xla-1.13
```
Cloud TPU リソースの環境変数を構成します。

注: TPU_IP_ADDRESS 変数は、ポート番号なしで Cloud TPU リソースを起動する際に特定される Cloud TPU の IP アドレスと同じである必要があります。これは、Google Cloud Console の [Compute Engine] > [TPU] に表示される内部 IP アドレスです。

データをダウンロードして準備する

このタスクで使用できる代替データセットについては、NearlineR ウェブサイトをご覧ください。このチュートリアルでは、前処理時間が最も短い dev-clean.tar.gz を使用します。

Wav2Vec2 には、必要な依存関係がいくつかあります。ここでは、それをインストールします。
```
(vm) $ pip install omegaconf hydra-core soundfile
(vm) $ sudo apt-get install libsndfile-dev
```

データセットをダウンロードします。

(vm) $ curl https://www.openslr.org/resources/12/dev-clean.tar.gz --output dev-clean.tar.gz

圧縮ファイルを展開します。ファイルは、LibriSpeech フォルダに保存されます。
```
(vm) $ tar xf dev-clean.tar.gz
```

最新の fairseq モデルをダウンロードしてインストールします。

(vm) $ git clone --recursive https://github.com/pytorch/fairseq.git
(vm) $ cd fairseq
(vm) $ pip install --editable .
(vm) $ cd -

データセットを準備します。このスクリプトでは、元データ（LibriSpeech/ の下）へのポインタを持つ manifest/ という名前のフォルダを作成します。
```
(vm) $ python fairseq/examples/wav2vec/wav2vec_manifest.py LibriSpeech/ --dest manifest/
```

トレーニングジョブを実行する

LibriSpeech データでモデルを実行します。このスクリプトの実行には、約 2 時間かかります。

(vm) $ fairseq-hydra-train \
task.data=${HOME}/manifest \
optimization.max_update=400000 \
dataset.batch_size=4 \
common.log_format=simple \
--config-dir fairseq/examples/wav2vec/config/pretraining   \
--config-name wav2vec2_large_librivox_tpu.yaml

クリーンアップ

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

Compute Engine インスタンスとの接続を切断していない場合は切断します。
```
(vm)$ exit
```
プロンプトが user@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。
Cloud Shell で Google Cloud CLI を使用して Compute Engine VM インスタンスと TPU を削除します。
```
$ gcloud compute tpus delete wav2vec2-tutorial --zone=us-central1-a
```

次のステップ

Cloud TPU Pod へのスケーリング

このチュートリアルの事前トレーニングタスクを強力な Cloud TPU Pod にスケールするには、Cloud TPU Pod での PyTorch モデルのトレーニングチュートリアルをご覧ください。

次のように PyTorch colabs を試す