Vertex AI Experimentsを利用したMNISTモデル学習の実験管理
今回はVertex AIのExperimentsを利用して、MNISTの学習家庭を記録してみました。Vertex AI Experimentsを利用することでモデルの実験評価ができ、どのモデルをデプロイするのが良いかを判断する材料になります。
Vertex AI Experimentsとは？
Vertex AI Experiments（以下、Experiments）とは、Google CloudのVertex AI上で提供されている実験管理ツールになります。Experimentsを利用することでモデルアーキテクチャからハイパーパラメータ、トレーニング結果まで幅広い情報を記録して分析することができ、最適なMLモデルの選定に役立ちます。
早速使ってみる
検証内容
今回はPytorchを利用してMNISTの分類モデルを実装し、その結果をExperimentsで記録し、モデルの性能を比較してみようと思います。なお実行環境はVertex AI Workbenchを利用しました。Workbenchの設定については全ての設定をデフォルトで設定しておきました（GPUは利用していないです）。また、今回の検証はExperimentsの使い勝手の検証なので、モデルの精度とかは分析しないです。
Python環境構築
今回はWorkbenchで実行したので、pipを使って以下のライブラリをインストールしました。
pip install torch torchvision tqdm google-cloud-aiplatform["autologging"] tensorboard
データ前処理の実装
今回対象とするデータはMNISTであり、torchvisionから簡単に取得できます。今回は最低限の処理に止めて実装します。
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
TRAIN_BATCH_SIZE = 1024
TEST_BATCH_SIZE = 1024
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=TRAIN_BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=TEST_BATCH_SIZE, shuffle=False)
モデルの定義
簡単なCNNを利用したモデルを定義します。実験で比較するためにCNNのフィルタ数をパラメータで変更できるようにします。
from torch import nn
class SimpleCNN(nn.Module):
def __init__(self, conv1_filters: int, conv2_filters: int):
super(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(1, conv1_filters, kernel_size=5)
self.conv2 = nn.Conv2d(conv1_filters, conv2_filters, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(conv2_filters * 16, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = nn.functional.relu(nn.functional.max_pool2d(self.conv1(x), 2))
x = nn.functional.relu(nn.functional.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, x.size(1) * x.size(2) * x.size(3))
x = nn.functional.relu(self.fc1(x))
x = nn.functional.dropout(x, training=self.training)
x = self.fc2(x)
return nn.functional.log_softmax(x, dim=1)
学習コードの実装
それではデータとモデルの実装を参照しながら、Experimentsに記録しつつ学習を進めるコードを実装します。
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
import os
import shutil
from itertools import product
from tqdm import tqdm
from uuid import uuid4
import os
import google.cloud.aiplatform as aip
from model import SimpleCNN
from prepare_data import train_loader, test_loader, TRAIN_BATCH_SIZE, TEST_BATCH_SIZE
PROJECT_ID = "..."
LOCATION = "..."
aip.init(project=PROJECT_ID, location=LOCATION)
EXPERIMENT_NAME = "torch-mnist-experiment"
aip.init(experiment=EXPERIMENT_NAME)
aip.autolog()
GCS_BUCKET_NAME = "torch-mnist-experiment"
CONV1_FILTERS = [3, 5, 7]
CONV2_FILTERS = [9, 11, 13]
LRS = [0.1, 0.01, 0.001]
for conv1_filters, conv2_filters, lr in tqdm(product(CONV1_FILTERS, CONV2_FILTERS, LRS)):
run_name = f"simple-cnn-{uuid4()}"
aip.start_run(run_name)
epochs = 5
aip.log_params({
"learning_rate": lr,
"epochs": epochs,
"optimizer": "Adam",
"model_architecture": "SimpleCNN",
"conv1_filters": conv1_filters,
"conv2_filters": conv2_filters,
"train_batch_size": TRAIN_BATCH_SIZE,
"test_batch_size": TEST_BATCH_SIZE,
"device": "cpu",
})
TENSORBOARD_LOG_DIR = f"./runs/{run_name}"
if os.path.exists(TENSORBOARD_LOG_DIR):
shutil.rmtree(TENSORBOARD_LOG_DIR)
writer = SummaryWriter(TENSORBOARD_LOG_DIR)
model = SimpleCNN(conv1_filters=conv1_filters, conv2_filters=conv2_filters)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
for epoch in tqdm(range(epochs), leave=False):
model.train()
running_loss = 0.0
for data, target in train_loader:
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
running_loss += loss.item()
avg_train_loss = running_loss / len(train_loader)
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in tqdm(test_loader, leave=False):
output = model(data)
test_loss += criterion(output, target).item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
avg_test_loss = test_loss / len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
aip.log_time_series_metrics({"avt_test_loss": avg_test_loss, "accuracy": accuracy})
writer.add_scalar('Loss/train', avg_train_loss, epoch)
writer.add_scalar('Loss/test', avg_test_loss, epoch)
writer.add_scalar('Accuracy', accuracy, epoch)
writer.close()
GCS_TB_LOG_PATH = f"gs://{GCS_BUCKET_NAME}/{EXPERIMENT_NAME}/{run_name}"
os.system(f"gsutil -m cp -r {TENSORBOARD_LOG_DIR}/* {GCS_TB_LOG_PATH}")
aip.end_run()
まずはプロジェクト設定およびExperimentsの設定を以下の部分で実施しています。Experimentsの実験名としてはtorch-mnist-experimentとして実験を作成しつつ、aip.autolog()にて実験のログの自動記録をONにしています。この機能を利用するためにはgoogle-cloud-aiplatform["autologging"]の利用が必要なのですが、内部的にはmlflowを利用しているようでした。
PROJECT_ID = "..."
LOCATION = "..."
aip.init(project=PROJECT_ID, location=LOCATION)
EXPERIMENT_NAME = "torch-mnist-experiment"
aip.init(experiment=EXPERIMENT_NAME)
aip.autolog()
次にパラメータですが、今回はconv1とconv2のフィルタ数、および学習率を3パターンずつ用意しました。なお、エポックについては5に固定しています。
CONV1_FILTERS = [3, 5, 7]
CONV2_FILTERS = [9, 11, 13]
LRS = [0.1, 0.01, 0.001]
for conv1_filters, conv2_filters, lr in tqdm(product(CONV1_FILTERS, CONV2_FILTERS, LRS)):
...
ループの中ではまずパラメータセットの記録をします。実験の中ではrunという単位で個別の実験が記録されますが、その実験内容のパラメータを設定しています。
aip.log_params({
"learning_rate": lr,
"epochs": epochs,
"optimizer": "Adam",
"model_architecture": "SimpleCNN",
"conv1_filters": conv1_filters,
"conv2_filters": conv2_filters,
"train_batch_size": TRAIN_BATCH_SIZE,
"test_batch_size": TEST_BATCH_SIZE,
"device": "cpu",
})
以下の実装ではTensorboardへの記録をできるように準備しています。
TENSORBOARD_LOG_DIR = f"./runs/{run_name}"
if os.path.exists(TENSORBOARD_LOG_DIR):
shutil.rmtree(TENSORBOARD_LOG_DIR)
writer = SummaryWriter(TENSORBOARD_LOG_DIR)
以下では学習の実行およびテストの実行をしており、ループごとに精度やロスを記録しています。
for epoch in tqdm(range(epochs), leave=False):
model.train()
running_loss = 0.0
for data, target in train_loader:
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
running_loss += loss.item()
avg_train_loss = running_loss / len(train_loader)
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in tqdm(test_loader, leave=False):
output = model(data)
test_loss += criterion(output, target).item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
avg_test_loss = test_loss / len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
aip.log_time_series_metrics({"avt_test_loss": avg_test_loss, "accuracy": accuracy})
writer.add_scalar('Loss/train', avg_train_loss, epoch)
writer.add_scalar('Loss/test', avg_test_loss, epoch)
writer.add_scalar('Accuracy', accuracy, epoch)
runが終わったタイミングでrunの結果をGCSに保存し、runの記録を終了させます。
writer.close()
GCS_TB_LOG_PATH = f"gs://{GCS_BUCKET_NAME}/{EXPERIMENT_NAME}/{run_name}"
os.system(f"gsutil -m cp -r {TENSORBOARD_LOG_DIR}/* {GCS_TB_LOG_PATH}")
aip.end_run()
実行結果の確認
こちらのコードを実行すると、Vertex AIの「モデルの開発」の「テスト」に以下のようにExperimentが追加されます。
表示されたExperimentを選択すると、その実験に関連づけられたrunの一覧が表示されます。また、表示選択したrunの時系列結果が画面上部に表示されます。
複数のrunについて結果を比較したい場合、比較対象のrunを以下のように選択します。
選択した状態で「比較」を選択すると以下のように比較をするための画面が表示されます。こちらの画面でパラメータや評価指標の比較などを実施することができます。
このほかに、Tensorboardを利用した実験管理もできます。今回はTensorboardと連携させているので、最初の画像にあるOpen Tensorboardをクリックすると新規タブでTensorboardが表示されます。
まとめ
今回はVertex AI Experimentsを使ってみました。普段実験管理する時はmlflowを使うことが多かったですが、Vertex AI上でMLOpsを構築する場合はExperimentsの採用が第一候補になると思います。今回は最低限のSDKの利用に止めていますが、今後はより複雑な実験管理も試してみようと思います。
Discussion