このチュートリアルでは、Cloud TPU で FairSeq の RoBERTa を事前トレーニングする方法を説明します。具体的には、公開されている wikitext-103 データセットでモデルを事前トレーニングする FairSeq のチュートリアルに従います。
目標
- PyTorch 環境を作成および構成する
- データセットを準備する
- トレーニング ジョブを実行する
- 出力結果を表示できることを確認する
費用
このチュートリアルでは、Google Cloud の課金対象となる以下のコンポーネントを使用します。- Compute Engine
- Cloud TPU
料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを出すことができます。新しい Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。
始める前に
このチュートリアルを開始する前に、Google Cloud プロジェクトが正しく設定されていることを確認します。
- Google アカウントにログインします。
Google アカウントをまだお持ちでない場合は、新しいアカウントを登録します。
-
Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。
-
Cloud プロジェクトに対して課金が有効になっていることを確認します。プロジェクトに対して課金が有効になっていることを確認する方法を学習する。
このチュートリアルでは、Google Cloud の課金対象となるコンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。不要な課金を回避するために、このチュートリアルを完了したら、作成したリソースを必ずクリーンアップしてください。
Compute Engine インスタンスを設定する
Cloud Shell ウィンドウを開きます。
プロジェクト ID の変数を作成します。
export PROJECT_ID=project-id
Cloud TPU を作成するプロジェクトを使用するように
gcloud
コマンドライン ツールを構成します。gcloud config set project ${PROJECT_ID}
このコマンドを新しい Cloud Shell VM で初めて実行すると、
Authorize Cloud Shell
ページが表示されます。ページの下部にある [Authorize
] をクリックして、gcloud
に認証情報を使用した GCP API の呼び出しを許可します。Cloud Shell から、このチュートリアルで必要となる Compute Engine リソースを起動します。
gcloud compute instances create roberta-tutorial \ --zone=us-central1-a \ --machine-type=n1-standard-16 \ --image-family=torch-xla \ --image-project=ml-images \ --boot-disk-size=200GB \ --scopes=https://www.googleapis.com/auth/cloud-platform
新しい Compute Engine インスタンスに接続します。
gcloud compute ssh roberta-tutorial --zone=us-central1-a
Cloud TPU リソースを起動する
Compute Engine 仮想マシンから、次のコマンドを使用して Cloud TPU リソースを起動します。
(vm) $ gcloud compute tpus create roberta-tutorial \ --zone=us-central1-a \ --network=default \ --version=pytorch-1.7 \ --accelerator-type=v3-8
Cloud TPU リソースの IP アドレスを識別します。
(vm) $ gcloud compute tpus list --zone=us-central1-a
PyTorch 環境を作成および構成する
conda
環境を開始します。(vm) $ conda activate torch-xla-1.7
Cloud TPU リソースの環境変数を構成します。
(vm) $ export TPU_IP_ADDRESS=ip-address
(vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
データを設定する
以下を実行して FairSeq をインストールします。
(vm) $ pip install --editable /usr/share/torch-xla-1.7/tpu-examples/deps/fairseq
モデルデータを格納する pytorch-tutorial-data ディレクトリを作成します。
(vm) $ mkdir $HOME/pytorch-tutorial-data (vm) $ cd $HOME/pytorch-tutorial-data
FairSeq RoBERTa の README「データの前処理」セクションの手順に従います。 データセットの準備には約 10 分かかります。
モデルのトレーニング
モデルをトレーニングするには、まず環境変数を設定します。
(vm) $ export TOTAL_UPDATES=125000 # Total number of training steps (vm) $ export WARMUP_UPDATES=10000 # Warmup the learning rate over this many updates (vm) $ export PEAK_LR=0.0005 # Peak learning rate, adjust as needed (vm) $ export TOKENS_PER_SAMPLE=512 # Max sequence length (vm) $ export UPDATE_FREQ=16 # Increase the batch size 16x
(vm) $ export DATA_DIR=${HOME}/pytorch-tutorial-data/data-bin/wikitext-103
次のスクリプトを実行します。
(vm) $ python /usr/share/torch-xla-1.7/tpu-examples/deps/fairseq/train.py $DATA_DIR \
--task masked_lm \
--criterion masked_lm \
--arch roberta_base \
--sample-break-mode complete \
--tokens-per-sample $TOKENS_PER_SAMPLE \
--optimizer adam \
--adam-betas '(0.9,0.98)' \
--adam-eps 1e-6 \
--clip-norm 0.0 \
--lr-scheduler polynomial_decay \
--lr $PEAK_LR \
--warmup-updates $WARMUP_UPDATES \
--total-num-update $TOTAL_UPDATES \
--dropout 0.1 \
--attention-dropout 0.1 \
--weight-decay 0.01 \
--update-freq $UPDATE_FREQ \
--max-update $TOTAL_UPDATES \
--log-format simple \
--valid-subset=valid \
--train-subset=train \
--num_cores=8 \
--metrics_debug \
--input_shapes 16x512 18x480 21x384 \
--save-dir=${HOME}/checkpoints \
--log_steps=30 \
--max-epoch=1 \
--skip-invalid-size-inputs-valid-test
トレーニング スクリプトは約 15 分間実行され、終了すると次のようなメッセージが表示されます。
saved checkpoint /home/user/checkpoints/checkpoint1.pt (epoch 1 @ 119 updates) (writing took 25.19265842437744 seconds) | done training in 923.8 seconds
出力結果を確認する
トレーニング ジョブが完了の後、モデル チェックポイントが次のディレクトリに置かれます。
$HOME/checkpoints
クリーンアップ
作成したリソースを使用した後、アカウントに不要な請求が発生しないようにクリーンアップを行います。
Compute Engine インスタンスとの接続を切断していない場合は切断します。
(vm) $ exit
プロンプトが
user@projectname
に変わります。これは、現在、Cloud Shell 内にいることを示しています。Cloud Shell で、
gcloud
コマンドライン ツールを使用して、Compute Engine インスタンスを削除します。$ gcloud compute instances delete roberta-tutorial --zone=us-central1-a
gcloud
コマンドライン ツールを使用して、Cloud TPU リソースを削除します。$ gcloud compute tpus delete roberta-tutorial --zone=us-central1-a
次のステップ
次のように PyTorch colabs を試す
- Cloud TPU での PyTorch のスタートガイド
- TPU 上で MNIST をトレーニングする
- Cifar10 データセットを使用して TPU 上で ResNet18 をトレーニングする
- Pretrained ResNet50 モデルを使用して推論する
- 高速なニューラルスタイル変換
- Fashion MNIST を使用した AlexNet のマルチコア トレーニング
- Fashion MNIST を使用した AlexNet のシングルコア トレーニング