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

1. 概要

このラボでは、Google Cloud での ML トレーニングワークフロー全体を説明します。モデルの構築には PyTorch を使用します。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 の [モデル] セクションに移動し、まだ有効になっていない場合は [有効にする] をクリックします。

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

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

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

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

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

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

これで準備は完了です。

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

ノートブックの最初のセルに次のインポートを追加し、セルを実行します。実行するには、上部メニューの右矢印ボタンを押すか、command+enter を押します。

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

ここでは PyTorch をインポートしていません。これは、Notebook インスタンスからではなく、AI Platform Training でトレーニングジョブを実行しているためです。

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

AI Platform Training でトレーニングジョブを実行するには、トレーニングコードを Notebooks インスタンスでローカルにパッケージ化し、ジョブのアセットを保存する 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()

このデータセットには 10 万行弱のデータが含まれています。ここでは、母親と父親の年齢、妊娠週数、母親の体重増加（ポンド単位）、新生児の性別（ブール値）の 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 テンソルに変換してから、モデルを定義します。モデルを構築するために、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: ジョブをモニタリングする

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

最初のトライアルが開始されると、各トライアルで選択されたハイパーパラメータ値が表示されます。

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

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

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

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

7. クリーンアップ

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

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

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