🔥

Aimシリーズ：OptunaとPytorch Lightningを組み合わせたMNIST実験管理

2025/08/23に公開

今回はAimで実験管理を行いつつ、OptunaとPytorch Lightningを使ってMNISTの分類をしてみました。ぜひ過去の以下の記事を参考にしてください。

早速実装

環境構築

uvを使って以下で環境を構築します。

uv init aim_optuna_lightning_mnist -p 3.12
cd aim_optuna_lightning_mnist
uv add aim lightning optuna torch torchvision

コードの実装

今回利用するコードは以下になります。

main.py

from aim import Run
import optuna
import torch
import os
from torch import optim, nn, utils, Tensor
from torchvision.datasets import MNIST
from torchvision import transforms
import lightning as L


def calculate_accuracy(p, label):
    p_arg = torch.argmax(p, dim=1)
    return torch.sum(label == p_arg)


class LitClassifier(L.LightningModule):
    def __init__(self, net):
        super().__init__()
        self.net = net
        self.criteria = nn.CrossEntropyLoss()
        self.train_loss = 0
        self.val_loss = 0
        self.train_acc = 0
        self.val_acc = 0

    def training_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), 28*28)
        out = self.net(x)
        loss = self.criteria(out, y)
        self.log("train_loss", loss)
        self.train_loss += loss.item()
        correct_num = calculate_accuracy(out, y)
        self.train_acc += correct_num.item()
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), 28*28)
        out = self.net(x)
        loss = self.criteria(out, y)
        self.log("val_loss", loss)
        self.val_loss += loss.item()
        correct_num = calculate_accuracy(out, y)
        self.val_acc += correct_num.item()
        return loss

    def configure_optimizers(self):
        optimizer = optim.SGD(self.net.parameters(), lr=1e-3)
        return optimizer


class Net(nn.Module):
    def __init__(self, layer1_size: int, layer2_size: int, layer3_size: int):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28*28, layer1_size)
        self.fc2 = nn.Linear(layer1_size, layer2_size)
        self.fc3 = nn.Linear(layer2_size, layer3_size)
        self.fc4 = nn.Linear(layer3_size, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        out = self.relu(self.fc1(x))
        out = self.relu(self.fc2(out))
        out = self.relu(self.fc3(out))
        return self.fc4(out)


def create_model(layer1_size, layer2_size, layer3_size):
  return LitClassifier(
        Net(layer1_size, layer2_size, layer3_size)
    )


transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, ), (0.5, ))]
)
train_dataset = MNIST(os.getcwd(), download=True, transform=transform, train=True)
val_dataset = MNIST(os.getcwd(), download=True, transform=transform, train=False)

def objective(trial: optuna.trial.Trial):
    run = Run()

    # We optimize the number of layers, hidden units in each layer and dropouts.
    layer1_size = trial.suggest_int("layer1_size", 500, 1_000)
    layer2_size = trial.suggest_int("layer2_size", 1_000, 2_000)
    layer3_size = trial.suggest_int("layer3_size", 500, 1_000)
    batch_size = trial.suggest_int("batch_size", 128, 1024)
    train_loader = utils.data.DataLoader(train_dataset, batch_size)
    val_loader = utils.data.DataLoader(val_dataset, batch_size)
    classifier = create_model(layer1_size, layer2_size, layer3_size)
    run["hparams"] = {"layer1_size": layer1_size, "layer2_size": layer2_size, "layer3_size": layer3_size, "batch_size": batch_size}

    class EpochCallback(L.Callback):
        def on_train_epoch_start(self, trainer, pl_module):
            pl_module.train_loss = 0
            pl_module.val_loss = 0
            pl_module.train_acc = 0
            pl_module.val_acc = 0

        def on_train_epoch_end(self, trainer, pl_module):
            run.track(pl_module.train_loss, name='loss', step=pl_module.current_epoch, context={ "subset":"train" })
            run.track(pl_module.val_loss, name='loss', step=pl_module.current_epoch, context={ "subset":"val" })
            run.track(pl_module.train_acc * 100 / len(train_dataset), name='acc', step=pl_module.current_epoch, context={ "subset":"train" })
            run.track(pl_module.val_acc * 100 / len(val_dataset), name='acc', step=pl_module.current_epoch, context={ "subset":"val" })


    trainer = L.Trainer(max_epochs=10, callbacks=[EpochCallback()])
    trainer.fit(model=classifier, train_dataloaders=train_loader, val_dataloaders=val_loader)
    trainer.logger.log_hyperparams({"layer1_size": layer1_size, "layer2_size": layer2_size, "layer3_size": layer3_size})
    return trainer.callback_metrics["train_loss"].item() + (layer1_size ** 2 + layer2_size ** 2 + layer3_size ** 2) ** 0.5

# train the model (hint: here are some helpful Trainer arguments for rapid idea iteration)

study = optuna.create_study()
study.optimize(objective, n_trials=30)

まずはLightningのTrainerを以下のように実装します。Aimで記録するにあたりtrainとvalでそれぞれlossとaccuracyを記録しています。

class LitClassifier(L.LightningModule):
    def __init__(self, net):
        super().__init__()
        self.net = net
        self.criteria = nn.CrossEntropyLoss()
        self.train_loss = 0
        self.val_loss = 0
        self.train_acc = 0
        self.val_acc = 0

    def training_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), 28*28)
        out = self.net(x)
        loss = self.criteria(out, y)
        self.log("train_loss", loss)
        self.train_loss += loss.item()
        correct_num = calculate_accuracy(out, y)
        self.train_acc += correct_num.item()
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), 28*28)
        out = self.net(x)
        loss = self.criteria(out, y)
        self.log("val_loss", loss)
        self.val_loss += loss.item()
        correct_num = calculate_accuracy(out, y)
        self.val_acc += correct_num.item()
        return loss

    def configure_optimizers(self):
        optimizer = optim.SGD(self.net.parameters(), lr=1e-3)
        return optimizer

次にネットワークではLinearのみを使って実装し、中間層3つのニューロン数をパラメータとします。

class Net(nn.Module):
    def __init__(self, layer1_size: int, layer2_size: int, layer3_size: int):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28*28, layer1_size)
        self.fc2 = nn.Linear(layer1_size, layer2_size)
        self.fc3 = nn.Linear(layer2_size, layer3_size)
        self.fc4 = nn.Linear(layer3_size, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        out = self.relu(self.fc1(x))
        out = self.relu(self.fc2(out))
        out = self.relu(self.fc3(out))
        return self.fc4(out)

今回取り扱うデータはtorchvisionからMNISTデータを利用します。今回はテストデータをvalデータとして扱います。

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, ), (0.5, ))]
)
train_dataset = MNIST(os.getcwd(), download=True, transform=transform, train=True)
val_dataset = MNIST(os.getcwd(), download=True, transform=transform, train=False)

最後にOptunaで管理される目的関数を実装します。今回のハイパーパラメータは中間層のニューロン数とバッチサイズとしています。また、Aimに記録するにあたり学習エポックの前後で記録するためにLightningのコールバック機能を利用します。EpochCallbackにてコールバックを実装します。on_train_epoch_startが学習エポックが開始する前に、on_train_epoch_endがエポック終了後に実行されるコールバックであり、それぞれメトリクスのリセットとメトリクスの記録をしています。

def objective(trial: optuna.trial.Trial):
    run = Run()

    # We optimize the number of layers, hidden units in each layer and dropouts.
    layer1_size = trial.suggest_int("layer1_size", 500, 1_000)
    layer2_size = trial.suggest_int("layer2_size", 1_000, 2_000)
    layer3_size = trial.suggest_int("layer3_size", 500, 1_000)
    batch_size = trial.suggest_int("batch_size", 128, 1024)
    train_loader = utils.data.DataLoader(train_dataset, batch_size)
    val_loader = utils.data.DataLoader(val_dataset, batch_size)
    classifier = create_model(layer1_size, layer2_size, layer3_size)
    run["hparams"] = {"layer1_size": layer1_size, "layer2_size": layer2_size, "layer3_size": layer3_size, "batch_size": batch_size}

    class EpochCallback(L.Callback):
        def on_train_epoch_start(self, trainer, pl_module):
            pl_module.train_loss = 0
            pl_module.val_loss = 0
            pl_module.train_acc = 0
            pl_module.val_acc = 0

        def on_train_epoch_end(self, trainer, pl_module):
            run.track(pl_module.train_loss, name='loss', step=pl_module.current_epoch, context={ "subset":"train" })
            run.track(pl_module.val_loss, name='loss', step=pl_module.current_epoch, context={ "subset":"val" })
            run.track(pl_module.train_acc * 100 / len(train_dataset), name='acc', step=pl_module.current_epoch, context={ "subset":"train" })
            run.track(pl_module.val_acc * 100 / len(val_dataset), name='acc', step=pl_module.current_epoch, context={ "subset":"val" })


    trainer = L.Trainer(max_epochs=10, callbacks=[EpochCallback()])
    trainer.fit(model=classifier, train_dataloaders=train_loader, val_dataloaders=val_loader)
    trainer.logger.log_hyperparams({"layer1_size": layer1_size, "layer2_size": layer2_size, "layer3_size": layer3_size})
    return trainer.callback_metrics["train_loss"].item() + (layer1_size ** 2 + layer2_size ** 2 + layer3_size ** 2) ** 0.5

学習の実行

学習の実行は以下のように実行します。

uv run main.py

実験記録の確認

Aimの実験管理をUIにてするため、以下のコマンドを実行します。

uv run aim up

実行後に127.0.0.1:43800にアクセスして、Run一覧に移動すると以下のように30通りのRunが確認できます。

全て選択してCompareからmetricsを選択すると以下のように指標比較ができます。

グラフにホバーすると以下のようにそのグラフの情報を確認できます。

結果としてはlayer1_size=729、layer2_size=1087, layer3_size=649の時、学習精度が78.7になり、70.0になりました。またAimは実行時のCPU使用率はメモリ使用率も取得でき、インフラの検討材料になります。

また、個別の実験結果を確認すると、実行時のログも自動記録されているので、どのような処理が行われたかを別で記録する必要がなく便利です。

まとめ

今回はAimを利用しながらOptunaやLightningを使ってモデル開発の実験管理をしてみました。実験管理だけでなくログやインフラの内容も一括管理できるのでとても便利だと思います。ぜひ皆さんも実験管理ツールとしてAimの仕様を検討してみてください。