はじめに
EVERSTEELで業務委託の機械学習エンジニアをしている高濱です！
EVERSTEELでは、鉄スクラップ画像のAI解析を行うプロダクト「鉄ナビ検収AI」を開発しています。
そのコア機能の一つとして物体検出モデルを構築しているのですが、モデルの保守性向上や改善サイクルの高速化のため、学習や実験管理にMetaflowを導入しました。
この記事ではその検証過程で得られた知見として、Metaflowを使ってAWS Batch上で簡単なMLモデルを動かし、MNIST分類を実行するところまでを紹介したいと思います。
Metaflow on AWS Batch を使うモチベーション
以前の記事（最新論文調査：大規模モデル時代における鉄スクラップ CVの現在地）でも紹介されていますが、EVERSTEELでは物体検出モデルを使った異物検出を行っています。
EVERSTEELではクラウド基盤としてAWSを使っており、これまでは開発者がAWS EC2インスタンスを建て、その中でモデル学習の実験などを行なっていました。
しかし、以下のような問題が発生しました。
- 単独のEC2インスタンスでは、複雑な実験を並列で回したい場合にリソースに限界がある。同じEC2環境を複数構築するのも面倒。
- 実験環境や実験設定の管理が俗人化していて把握しにくく、予期せぬバグの原因になる。
これを解消するため、既に社内の別タスクで導入されていたAWS Batchを使うことを検討しました。
AWS Batchは、バッチ処理をフルマネージドで効率的に実行・管理できるサービスです。まとまった計算ジョブの定期実行などにも適しています。
必要に応じてEC2インスタンスなどを起動してバッチ処理を行えるため、事前に準備した環境で複数の実験を並列で回すことができることが期待されます。
ただ、AWS Batchを使う場合、通常は実行用のDocker ImageをBuild & Pushする必要があります。機械学習モデルを探索的に実験する場合など、手元で色々実装を変えつつジョブを回したい場合、毎回Build & Pushするのは少し面倒です。
（回避する方法は色々あるとは思いますが）この面倒を解消するための有力な方法として、今回はMetaflowを導入しました。
MetaflowはNetflixによって開発された、ジョブのワークフローシステムです。
Stepを定義してシステマティックにワークフローを定義できるほか、各種クラウドや外部ツールなどの連携も簡単にできるようになっており、利便性、保守性、再現性などに優れたツールとなっています。
MetaflowにはAWS Batchとの連携機能も用意されており、AWS Batchを使う場合は、いくつかの設定を書いてRunすることで、Docker Imageをbuildすることなく手元のコードをそのままAWS Batchで動かすことができます（便利！）
ワークフロー管理ツールという点では類似ツールは他にもありますが（Airflow, Prefect, Kubeflowなど）、今回は「既存のコードをなるべく変更することなく」「AWSの各種リソースと連携しつつ」「探索的な実験を行う用途で」動かしたいという目的だったため、それに適していそうなMetaflow on AWS Batchを検証しました。
Metaflow on AWS BatchでHello worldしてみる
では実際にMetaflowを使ってAWS Batchを動かしてみましょう。まずは単純なHello Worldを実行してみます。
なお、AWSへのログインなどは既に設定済みであることを想定しています。実行に必要な権限があれば、手元のローカルPC環境から実行しても、他のEC2インスタンスなどから実行しても構いません。
Install
公式のガイダンスに従ってMetaflowをinstallします
pip install metaflow
AWS Batch Job Queueを設定する
AWS Batchを使う際は、まずJob Definitionを作成し、その定義に沿ったJobをJob Queueに投げることで実行します。
Job DefinitionはMetaflowがジョブを投げるときに自動で定義してくれるので、手動で定義する必要はありません。
Job Queueは事前に設定しておく必要があります。詳細は公式ドキュメントなどを読んでいただければと思いますが、Job Queueを作成するためにはまずコンピューティング環境を設定する必要があります。
今回コンピューティング環境は、今後GPUを使うことを考えてtype=EC2, Managedで作成しました。
Job Queueもそれに合わせてtype=EC2で作成しました。名前はmetaflow-test-queueとしています。
AWS S3 Data Rootを設定する
Metaflowを使うと、ワークフローの実行ログをAWS S3に書き出すことができます。
今回はs3://es-metaflow-test/hello-world ディレクトリを作成しました。
Metaflow Configの設定
Metaflowを実行する際の環境変数は実行時に都度与えることもできますが、共通の設定は ~/.metaflowconfig/config.json に置いておくことで実行時に反映されます。
ここまでの設定を踏まえて以下のようなconfigを作成しました。
{
"METAFLOW_DEFAULT_DATASTORE": "s3",
"METAFLOW_DATASTORE_SYSROOT_S3": "s3://metaflow-test/hello-world",
"METAFLOW_DATATOOLS_S3ROOT": "s3://metaflow-test/hello-world/data",
"METAFLOW_ECS_S3_ACCESS_IAM_ROLE": "arn:aws:iam::XXX:role/ECSJobRole",
"METAFLOW_BATCH_JOB_QUEUE": "metaflow-test-queue",
}
Hello Worldするためのコードを書く
Hello Worldのコードは公式チュートリアルにあるものをベースに使います。
具体的には以下をmetaflow_helloworld.pyとして作成しました。
from metaflow import FlowSpec, batch, step
class HelloFlow(FlowSpec):
@step
def start(self):
print("HelloFlow is starting.")
self.next(self.hello)
@batch(cpu=1, queue="metaflow-test-queue")
@step
def hello(self):
print("Metaflow says: Hi!")
self.next(self.end)
@step
def end(self):
print("HelloFlow is all done.")
if __name__ == "__main__":
HelloFlow()
何をやっているかはコードを見ればなんとなく察してもらえるのではないかと思いますが、FlowSpecを継承したクラスを作成し、その中で@stepデコレーターをつけて関数を定義することで順に実行されます。startとendの2つの関数は必ず設定する必要があります（Creating Flows）
AWS Batchで動かしたい関数には@batchデコレーターを追加します。@batchデコレーターではcpu, gpu, memory, queueなどを個別に設定することができるため、stepごとに必要なリソース量を調整することができます。@batchを指定したコードはAWS Batchのジョブキューに送られ、つけない場合はローカルで実行されます。
ちなみに全部まとめて共通のbatch設定で実行するやり方もありますし、実行時にオプションとして渡す方法もあります。
Runする
以下を実行します
python metaflow_helloworld.py run
@batchをつけた関数を実行すると、AWS Batchで実行した場合と同じく、Statusが SUBMITTED → RUNNABLE → STARTING → RUNNING と変化していきます。
実行中のログストリームを見ると、startとendはローカルで実行され、hello関数だけBatchで実行されていることがわかります。
Batchの実行ログは、AWS BatchのUIからも確認できます。ジョブ一覧から指定したジョブキューで検索すると、実行結果の詳細やログを確認できます。
正常に終了したら、指定したS3 path s3://es-metaflow-test/hello-world にワークフローのログが書き出されているのを確認してください。
クラス名と同じHelloFlowフォルダが作成され、各stepの様々なログが保存されています。
Metaflow on AWS BatchでMNISTを学習する
Hello Worldができたので、続いてMNIST分類モデルを学習してみましょう
機械学習の実験管理ツールであるMLFlowと連携したり、学習済みモデルファイルをS3に保存したりすることも目指します（MLFlowの使い方についてはここでは触れません）。
Docker imageを登録する
MLFlowやs3と連携するためには、mlflow, boto3などのライブラリを追加でインストールする必要があります。
@pypi デコレーターなどを使って適宜インストールすることもできますが、今後さらに拡張性の高い環境構築が必要になるケースも見据えてDocker imageを使ってみます。
まず以下のような単純なDockerfileを作成します。
FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
WORKDIR /app
RUN pip install mlflow==2.5.0 boto3==1.28.85
作成したDocker imageは、ECRに登録します。
まずECRにリポジトリを作ります。名前はmetaflow-test-mnistとしました。
リポジトリを開いて「プッシュコマンドを表示」を見ると、docker imageのbuildからpushまでの手順が書いてあるので、それに従えばPushできると思います。
image nameはmetaflow-test-mnist:latestとしました。
なお、最初の課題感として「Docker Imageのbuildが面倒」というのがあったので、「結局Docker buildしてるじゃないか」と思った方もいるかもしれません。ただ、ここで構築したのはあくまで依存ライブラリなどのベース環境であり、一度pushすれば依存ライブラリなどが変わらない限り変更は必要ありません。
それも面倒な場合は@pypiデコレーターを使うか、あるいはパッケージ管理ツールであるuvとの連携も可能なようです。
Secretを登録する
MLFlowと連携する場合、username, passwordなどを環境変数として設定する必要がある場合があると思います。
このような場合は、Secrets Managerに登録することでMetaflowから利用することができます
今回はmlflowの名前でsecretを作り、key/value pairとしてMLFLOW_TRACKING_USERNAME, MLFLOW_TRACKING_PASSWORDを保存しました。
config.jsonを更新する
Hello Worldで作ったconfigに設定を追加します
- ECRのdefault resistoryを登録
- Secrets Managerを登録
{
"METAFLOW_DEFAULT_DATASTORE": "s3",
"METAFLOW_DATASTORE_SYSROOT_S3": "s3://metaflow-test/mnist",
"METAFLOW_DATATOOLS_S3ROOT": "s3://metaflow-test/mnist/data",
"METAFLOW_ECS_S3_ACCESS_IAM_ROLE": "arn:aws:iam::XXX:role/ECSJobRole",
"METAFLOW_BATCH_JOB_QUEUE": "metaflow-test-queue",
"METAFLOW_DEFAULT_CONTAINER_REGISTRY": "XXX.dkr.ecr.ap-northeast-1.amazonaws.com/",
"METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE": "aws-secrets-manager",
}
学習コードを書く
ここまでで準備が整ったので、学習用のコードを書いていきます。
コードはこちら！
import torch
import torch.nn as nn
from metaflow import FlowSpec, Parameter, batch, secrets, step
IMAGE = "metaflow-test-mnist:latest"
# モデルの定義 適当な数層のNN
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout1 = nn.Dropout(0.25)
self.dropout2 = nn.Dropout(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.conv1(x)
x = nn.functional.relu(x)
x = self.conv2(x)
x = nn.functional.relu(x)
x = nn.functional.max_pool2d(x, 2)
x = self.dropout1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
return nn.functional.log_softmax(x, dim=1)
class MNISTFlow(FlowSpec):
# 外部からパラメーターを入力することができる
# https://docs.metaflow.org/metaflow/basics#how-to-define-parameters-for-flows
epochs = Parameter("epochs", default=10, type=int, help="Number of training epochs")
batch_size = Parameter("batch_size", default=128, type=int, help="Mini batch size")
learning_rate = Parameter("learning_rate", default=0.005, type=float, help="Learning rate for training")
s3_model_path = Parameter(
"s3_model_path", default="s3://metaflow-test/mnist/mnist_model.pth", help="S3 path to save the model"
)
# その他のパラメーター
local_model_path = "./mnist_model.pth"
mlflow_uri = "https://mlflow.XXX.com/"
mlflow_experiment = "metaflow-test-mnist"
@batch(gpu=1, cpu=1, image=IMAGE)
@step
def start(self):
import torchvision
import torchvision.transforms as transforms
# データセットの準備
# 実行インスタンス内の./dataにデータセットをダウンロードする
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
self.train_dataset = torchvision.datasets.MNIST(root="./data", train=True, download=True, transform=transform)
self.test_dataset = torchvision.datasets.MNIST(root="./data", train=False, download=True, transform=transform)
print("MetaFlow is starting.")
self.next(self.train)
@secrets(sources=["eversteel-mlflow"]) # secretを指定
@batch(gpu=1, cpu=1, image=IMAGE) # gpuとdocker imageを指定
@step
def train(self):
import boto3
import mlflow
import torch
import torch.nn as nn
import torch.optim as optim
mlflow.set_tracking_uri(self.mlflow_uri)
mlflow.set_experiment(self.mlflow_experiment)
s3_client = boto3.client("s3")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_loader = torch.utils.data.DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)
model = Net().to(self.device)
optimizer = optim.Adam(model.parameters(), lr=self.learning_rate)
criterion = nn.CrossEntropyLoss()
# mlflowを開始して学習
with mlflow.start_run() as run:
self.run_id = run.info.run_id
print(f"MLFlow started! run_id: {self.run_id}")
model.train()
for epoch in range(self.epochs):
total_loss = 0
for data, target in train_loader:
data, target = data.to(self.device), target.to(self.device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(train_loader)
print(f"Epoch {epoch+1}/{self.epochs}, Loss: {avg_loss}")
mlflow.log_metric("train_loss", avg_loss, step=epoch)
# モデルをs3へ保存
torch.save(model.state_dict(), self.local_model_path)
bucket, key = self.s3_model_path.replace("s3://", "").split("/", 1)
s3_client.upload_file(self.local_model_path, bucket, key)
print(f"Model saved to {self.s3_model_path}")
self.next(self.evaluate)
@secrets(sources=["eversteel-mlflow"])
@batch(gpu=1, cpu=1, image=IMAGE)
@step
def evaluate(self):
import boto3
import mlflow
import torch
import torch.nn as nn
mlflow.set_tracking_uri(self.mlflow_uri)
mlflow.set_experiment(self.mlflow_experiment)
s3_client = boto3.client("s3")
bucket, key = self.s3_model_path.replace("s3://", "").split("/", 1)
s3_client.download_file(bucket, key, self.local_model_path)
model = Net().to(self.device)
model.load_state_dict(torch.load(self.local_model_path))
test_loader = torch.utils.data.DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False)
criterion = nn.CrossEntropyLoss()
with mlflow.start_run(run_id=self.run_id):
model.eval()
correct = 0
total = 0
test_loss = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(self.device), target.to(self.device)
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()
total += target.size(0)
test_loss /= len(test_loader)
accuracy = 100.0 * correct / total
print(f"Test Loss: {test_loss}, Accuracy: {accuracy}%")
mlflow.log_metric("test_loss", test_loss)
mlflow.log_metric("accuracy", accuracy)
self.accuracy = accuracy
self.next(self.end)
@step
def end(self):
"""End the flow."""
print(f"Training completed. Final accuracy: {self.accuracy}%")
if __name__ == "__main__":
MNISTFlow()
ポイント
-
Parameter()を使って実行時にパラメーターを渡すことができる -
@batchの中で、gpuとcustom imageを指定する -
＠secretsで、登録したmlflowの情報を読み込む - コード内でmlflow, boto3をimportして連携する
- Batch環境側でimportしたいライブラリは、ファイルの先頭ではなく関数の中でimportする
Runする
以下を実行します！
python metaflow_mnist.py run
こちらも実行が完了すると、S3やAWS Batchから実行ログを確認できると思います。
さらに、学習のメトリクスはコード内で指定したMLFlow path (experiment="metaflow-test-mnist") に保存されます
また、学習済みモデルはコード内で指定したs3 path (s3://metaflow-test/mnist/mnist_model.pth) に保存されます
10 epochほど学習して、テストデータで98%程度の精度が出ました。
MLFlowのUIから、学習ロスが低下していることも確認できました。
※失敗するケース
- RUNNABLEのまま動かない場合はメモリーが不足している可能性があります（参考）
- STARTINGのまま動かない場合は、Docker Imageが適切にpullできていないなどの問題が考えられます。
使ってみた感想
- AWS Batchを使う上で、docker imageのbuildが不要であったり、job definitionの設定が不要なのは便利でした
- 以前より簡単に実験を並列実行できるようになりました
- 今回使った以外にも、複雑な実行フローを定義したり、他ツールとの連携など幅広い機能が用意されており、拡張性が高そうです
- 実行が終了したら勝手にインスタンスが終了するのでお財布にも優しいです
- Metaflowの記法など学習コストがやや高い印象はあります
- ジョブの実行が開始するまで毎回数分ほど時間がかかるので、数分程度で終わるような短時間の実験にはあまり向いていないかもしれません
- 探索的な実験には便利ですが、本番運用に使ってはいないので向いているかどうかはわかりません
まとめ
この記事では、高速なMLモデルの実験サイクルを回すために活用しているMetaflow on AWS Batchの導入について紹介しました。
ここからさらにプロダクション向けの物体検出モデルを動かす際には、いくつかの機能追加や修正が必要だったので、その試行錯誤あれこれもいつか続編として記事にできればと思っています。
EVERSTEELでは、鉄スクラップ業界におけるML技術の実装を目指しながら、安全性、再現性、保守性の高いパイプライン構築を目指して日々実験や開発を行なっています。
我こそはスーパー機械学習パイプラインを作りたい！という方は、是非以下のリンクからさらにEVERSTEELについて知っていただければと思います！
Discussion