Cloud AI Platform での PyTorch モデルのトレーニングとハイパーパラメータ チューニング

1. 概要

このラボでは、PyTorch を使用してモデルを構築し、Google Cloud で ML トレーニング ワークフロー全体を確認します。ここでは、Cloud AI Platform Notebooks 環境で、AI Platform Training 上で実行できるようにトレーニング ジョブをパッケージ化し、ハイパーパラメータをチューニングする方法を学びます。

学習内容

次の方法を学習します。

  • AI Platform Notebooks インスタンスを作成する
  • PyTorch モデルを作成する
  • AI Platform Training でハイパーパラメータ チューニングを使用してモデルをトレーニングする

このラボを Google Cloud で実行するための総費用は約 $1 です。

2. 環境を設定する

この Codelab を実行するには、課金が有効になっている Google Cloud Platform プロジェクトが必要です。プロジェクトを作成するには、こちらの手順に沿って操作してください。

ステップ 1: Cloud AI Platform Models API を有効にする

まだ有効になっていない場合は、Cloud コンソールの [AI Platform モデル] セクションに移動して [有効にする] をクリックします。

d0d38662851c6af3.png

ステップ 2: Compute Engine API を有効にする

まだ有効になっていない場合は、[Compute Engine] に移動して [有効にする] を選択します。これはノートブック インスタンスを作成するために必要です。

ステップ 3: AI Platform Notebooks インスタンスを作成する

Cloud コンソールの [AI Platform Notebooks] セクションに移動し、[新しいインスタンス] をクリックします。次に、最新の PyTorch インスタンス タイプを選択します(GPU なし)。

892b7588f940d145.png

デフォルトのオプションを使用するか、必要に応じてカスタム名を付け、[作成] をクリックします。インスタンスが作成されたら、[JupyterLab を開く] を選択します。

63d2cf44801c2df5.png

次に、ランチャーから Python 3 ノートブック インスタンスを開きます。

de4c86c6c7f9438f.png

これで準備は完了です。

ステップ 5: Python パッケージをインポートする

ノートブックの最初のセルに次の import を追加してセルを実行します。実行するには、上部のメニューで右矢印ボタンを押すか、Command+Enter キーを押します。

import datetime
import numpy as np
import os
import pandas as pd
import time

ここでは PyTorch をインポートしていないことがわかります。これは、ノートブック インスタンスからではなく、AI Platform Training でトレーニング ジョブを実行するためです。

3. トレーニング ジョブのパッケージを作成する

AI Platform Training でトレーニング ジョブを実行するには、ノートブック インスタンスにローカルにパッケージ化したトレーニング コードと、ジョブのアセットを保存する Cloud Storage バケットが必要です。まず、ストレージ バケットを作成します。すでにアカウントがある場合は、この手順をスキップできます。

ステップ 1: モデル用の Cloud Storage バケットを作成する

まず、この Codelab の残りの部分で使用する環境変数を定義しましょう。以下の値に、Google Cloud プロジェクトの名前と作成する Cloud Storage バケットの名前(グローバルに一意である必要があります)を入力します。

# Update these to your own GCP project, model, and version names
GCP_PROJECT = 'your-gcp-project'
BOCKET_URL = 'gs://storage_bucket_name'

これで、ストレージ バケットを作成する準備が整いました。トレーニング ジョブを開始する際に参照します。

ノートブック内から次の gsutil コマンドを実行して、バケットを作成します。

!gsutil mb $BUCKET_URL

ステップ 2: Python パッケージの初期ファイルを作成する

AI Platform でトレーニング ジョブを実行するには、コードを Python パッケージとして構成する必要があります。これは、外部パッケージの依存関係を指定するルート ディレクトリ内の setup.py ファイル、パッケージの名前を含むサブディレクトリ(ここでは trainer/)、このサブディレクトリ内の空の __init__.py ファイルで構成されます。

まず、setup.py ファイルを作成します。iPython の %%writefile マジックを使用して、ファイルをインスタンスに保存します。ここでは、トレーニング コードで使用する 3 つの外部ライブラリ(PyTorch、Scikit-learn、Pandas)を指定しています。

%%writefile setup.py
from setuptools import find_packages
from setuptools import setup

REQUIRED_PACKAGES = ['torch>=1.5', 'scikit-learn>=0.20', 'pandas>=1.0']

setup(
    name='trainer',
    version='0.1',
    install_requires=REQUIRED_PACKAGES,
    packages=find_packages(),
    include_package_data=True,
    description='My training application package.'
)

次に、trainer/ ディレクトリとその中に空の init.py ファイルを作成します。Python はこのファイルを使用して、これがパッケージであることを認識します。

!mkdir trainer
!touch trainer/__init__.py

これで、トレーニング ジョブの作成を開始する準備が整いました。

4. データセットをプレビューする

このラボでは、モデルのトレーニングに使用するツールに焦点を当てていますが、モデルのトレーニングに使用するデータセットを簡単に確認しておきましょう。ここでは、BigQuery で利用可能な出生率データセットを使用します。これには、数十年にわたる米国の出生データが含まれます。このデータセットのいくつかの列を使用して、新生児の出生時体重を予測します。元のデータセットは非常に大きいので、Cloud Storage バケットで用意されているデータセットのサブセットを使用します。

ステップ 1: BigQuery の出生率データセットをダウンロードする

Cloud Storage で利用可能にしたバージョンのデータセットを Pandas DataFrame にダウンロードしてプレビューしてみましょう。

natality = pd.read_csv('https://storage.googleapis.com/ml-design-patterns/natality.csv')
natality.head()

このデータセットの行数は 100,000 を少し下回っています。新生児の出生時体重を予測するために、母親と父親の年齢、妊娠週数、母親の体重増加(ポンド単位)、赤ちゃんの性別をブール値で表した 5 つの特徴量を使用します。

5. ハイパーパラメータを調整してトレーニング ジョブを定義する

先ほど作成した trainer/ サブディレクトリ内の model.py というファイルにトレーニング スクリプトを記述します。トレーニング ジョブは AI Platform Training 上で実行されます。また、AI Platform のハイパーパラメータ調整サービスも利用して、ベイズ最適化を利用してモデルに最適なハイパーパラメータを見つけます。

ステップ 1: トレーニング スクリプトを作成する

まず、トレーニング スクリプトを使用して Python ファイルを作成します。次に、何が起こっているのかを詳しく見ていきます。この %%writefile コマンドを実行すると、モデルコードがローカルの Python ファイルに書き込まれます。

%%writefile trainer/model.py
import argparse
import hypertune
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim

from sklearn.utils import shuffle
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import normalize

def get_args():
    """Argument parser.
    Returns:
        Dictionary of arguments.
    """
    parser = argparse.ArgumentParser(description='PyTorch MNIST')
    parser.add_argument('--job-dir',  # handled automatically by AI Platform
                        help='GCS location to write checkpoints and export ' \
                             'models')
    parser.add_argument('--lr',  # Specified in the config file
                        type=float,
                        default=0.01,
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum',  # Specified in the config file
                        type=float,
                        default=0.5,
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--hidden-layer-size',  # Specified in the config file
                        type=int,
                        default=8,
                        help='hidden layer size')
    args = parser.parse_args()
    return args

def train_model(args):
    # Get the data
    natality = pd.read_csv('https://storage.googleapis.com/ml-design-patterns/natality.csv')
    natality = natality.dropna()
    natality = shuffle(natality, random_state = 2)
    natality.head()

    natality_labels = natality['weight_pounds']
    natality = natality.drop(columns=['weight_pounds'])


    train_size = int(len(natality) * 0.8)
    traindata_natality = natality[:train_size]
    trainlabels_natality = natality_labels[:train_size]

    testdata_natality = natality[train_size:]
    testlabels_natality = natality_labels[train_size:]

    # Normalize and convert to PT tensors
    normalized_train = normalize(np.array(traindata_natality.values), axis=0)
    normalized_test = normalize(np.array(testdata_natality.values), axis=0)

    train_x = torch.Tensor(normalized_train)
    train_y = torch.Tensor(np.array(trainlabels_natality))

    test_x = torch.Tensor(normalized_test)
    test_y = torch.Tensor(np.array(testlabels_natality))

    # Define our data loaders
    train_dataset = torch.utils.data.TensorDataset(train_x, train_y)
    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)

    test_dataset = torch.utils.data.TensorDataset(test_x, test_y)
    test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False)

    # Define the model, while tuning the size of our hidden layer
    model = nn.Sequential(nn.Linear(len(train_x[0]), args.hidden_layer_size),
                          nn.ReLU(),
                          nn.Linear(args.hidden_layer_size, 1))
    criterion = nn.MSELoss()

    # Tune hyperparameters in our optimizer
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
    epochs = 20
    for e in range(epochs):
        for batch_id, (data, label) in enumerate(train_dataloader):
            optimizer.zero_grad()
            y_pred = model(data)
            label = label.view(-1,1)
            loss = criterion(y_pred, label)
            
            loss.backward()
            optimizer.step()


    val_mse = 0
    num_batches = 0
    # Evaluate accuracy on our test set
    with torch.no_grad():
        for i, (data, label) in enumerate(test_dataloader):
            num_batches += 1
            y_pred = model(data)
            mse = criterion(y_pred, label.view(-1,1))
            val_mse += mse.item()


    avg_val_mse = (val_mse / num_batches)

    # Report the metric we're optimizing for to AI Platform's HyperTune service
    # In this example, we're mimizing error on our test set
    hpt = hypertune.HyperTune()
    hpt.report_hyperparameter_tuning_metric(
        hyperparameter_metric_tag='val_mse',
        metric_value=avg_val_mse,
        global_step=epochs        
    )

def main():
    args = get_args()
    print('in main', args)
    train_model(args)

if __name__ == '__main__':
    main()

トレーニング ジョブは、作業の大部分を行う 2 つの機能で構成されます。

  • get_args(): トレーニング ジョブの作成時に渡すコマンドライン引数と、AI Platform に最適化させるハイパーパラメータを解析します。この例の引数のリストには、最適化するハイパーパラメータ(モデルの学習率、モメンタム、隠れ層のニューロン数)のみが含まれています。
  • train_model(): ここでは、データを Pandas DataFrame にダウンロードして正規化し、PyTorch Tensor に変換してからモデルを定義します。モデルの構築には、PyTorch nn.Sequential API を使用します。これにより、モデルをレイヤのスタックとして定義できます。
model = nn.Sequential(nn.Linear(len(train_x[0]), args.hidden_layer_size),
                      nn.ReLU(),
                      nn.Linear(args.hidden_layer_size, 1))

ここでは、モデルの隠れ層のサイズをハードコードするのではなく、これをハイパーパラメータにして AI Platform が調整するようにしています。詳しくは次のセクションで説明します。

ステップ 2: AI Platform のハイパーパラメータ チューニング サービスを使用する

異なるハイパーパラメータ値を手動で試し、毎回モデルを再トレーニングする代わりに、Cloud AI Platform のハイパーパラメータ最適化サービスを使用します。ハイパーパラメータ引数を使用してトレーニング ジョブを設定すると、AI Platform はベイジアン最適化を使用して、指定したハイパーパラメータの理想的な値を見つけます。

ハイパーパラメータ チューニングでは、1 回のトライアルで、ハイパーパラメータ値の特定の組み合わせを使用してモデルのトレーニングを 1 回実施します。実行するトライアルの数に応じて、AI Platform は完了したトライアルの結果を使用して、将来のトライアル用に選択するハイパーパラメータを最適化します。ハイパーパラメータ調整を構成するには、トレーニング ジョブを開始するときに、最適化する各ハイパーパラメータに関するデータを含む構成ファイルを渡す必要があります。

次に、その構成ファイルをローカルに作成します。

%%writefile config.yaml
trainingInput:
  hyperparameters:
    goal: MINIMIZE
    maxTrials: 10
    maxParallelTrials: 5
    hyperparameterMetricTag: val_mse
    enableTrialEarlyStopping: TRUE
    params:
    - parameterName: lr
      type: DOUBLE
      minValue: 0.0001
      maxValue: 0.1
      scaleType: UNIT_LINEAR_SCALE
    - parameterName: momentum
      type: DOUBLE
      minValue: 0.0
      maxValue: 1.0
      scaleType: UNIT_LINEAR_SCALE
    - parameterName: hidden-layer-size
      type: INTEGER
      minValue: 8
      maxValue: 32
      scaleType: UNIT_LINEAR_SCALE

ハイパーパラメータごとに、型、検索する値の範囲、さまざまなトライアル間で値を増加させるスケールを指定します。

ジョブの開始時に、最適化する指標も指定します。上記の train_model() 関数の最後に、トライアルが完了するたびにこの指標を AI Platform に報告しています。ここでは、モデルの平均二乗誤差を最小化しているため、モデルの平均二乗誤差が最も小さいハイパーパラメータを使用します。この指標の名前(val_mse)は、トライアル終了時に report_hyperparameter_tuning_metric() を呼び出したときにレポートに使用する名前と一致します。

6. AI Platform でトレーニング ジョブを実行する

このセクションでは、AI Platform でハイパーパラメータを調整しながら、モデルのトレーニング ジョブを開始します。

ステップ 1: 環境変数を定義する

まず、トレーニング ジョブを開始するために使用する環境変数を定義します。別のリージョンでジョブを実行する場合は、以下の REGION 変数を更新します。

MAIN_TRAINER_MODULE = "trainer.model"
TRAIN_DIR = os.getcwd() + '/trainer'
JOB_DIR = BUCKET_URL + '/output'
REGION = "us-central1"

AI Platform の各トレーニング ジョブには一意の名前を付ける必要があります。次のコマンドを実行して、タイムスタンプを使用してジョブ名の変数を定義します。

timestamp = str(datetime.datetime.now().time())
JOB_NAME = 'caip_training_' + str(int(time.time()))

ステップ 2: トレーニング ジョブを開始する

Google Cloud CLI の gcloud を使用してトレーニング ジョブを作成します。上で定義した変数を参照して、このコマンドをノートブックで直接実行できます。

!gcloud ai-platform jobs submit training $JOB_NAME \
        --scale-tier basic \
        --package-path $TRAIN_DIR \
        --module-name $MAIN_TRAINER_MODULE \
        --job-dir $JOB_DIR \
        --region $REGION \
        --runtime-version 2.1 \
        --python-version 3.7 \
        --config config.yaml

ジョブが正しく作成された場合は、AI Platform コンソールの [ジョブ] セクションに移動してログをモニタリングします。

ステップ 3: ジョブをモニタリングする

コンソールの [ジョブ] セクションで、開始したジョブをクリックして詳細を表示します。

c184167641bb7ed7.png

最初のトライアルを開始すると、各トライアルに選択されたハイパーパラメータ値を確認できます。

787c053ef9110e6b.png

トライアルが完了すると、最適化指標の結果の値(この場合は val_mse)がここに記録されます。ジョブの実行には 15 ~ 20 分かかります。ジョブが完了すると、ダッシュボードは次のようになります(正確な値は異なります)。

47ef6b9b4ecb532c.png

潜在的な問題をデバッグし、ジョブを詳細にモニタリングするには、ジョブの詳細ページで [ログを表示] をクリックします。

18c32dcd36351930.png

モデル トレーニング コード内のすべての print() ステートメントがここに表示されます。問題が発生した場合は、print ステートメントを追加して新しいトレーニング ジョブを開始してみてください。

トレーニング ジョブが完了したら、val_mse が最も低いハイパーパラメータを見つけます。これらを使用して、モデルの最終バージョンをトレーニングしてエクスポートすることも、追加のハイパーパラメータ調整トライアルを含む別のトレーニング ジョブを開始するためのガイダンスとして使用することもできます。

7. クリーンアップ

このノートブックを引き続き使用する場合は、未使用時にオフにすることをおすすめします。Cloud コンソールの Notebooks UI で、ノートブックを選択して [停止] を選択します。

879147427150b6c7.png

このラボで作成したすべてのリソースを削除する場合は、ノートブック インスタンスを停止するのではなく削除します。

Cloud コンソールのナビゲーション メニューで [ストレージ] に移動し、モデルアセットを保存するために作成した両方のバケットを削除します。