🙌

OptunaとPyTorch Lightningを組み合わせてモデルの最適化をしてみた

2025/08/15に公開

今回はOptunaとPyTorch Lightningを組み合わせてモデルの最適化をしてみようと思います。なお、Optunaは昨日紹介した記事を、Pytorch Lightningについてはこちらの記事をベースに進めようと思います。

それでは早速やってみる！

環境構築

uvを使って環境構築していきます。

uv init optuna_pytorch_lightning -p 3.12
cd optuna_pytorch_lightning
uv add lightning optuna torch torchvision

ソースコードの実装

それでは次にソースコードを実装していきます。今回はエンコーダとデコーダの各レイヤー数を最適化するようにしてみます。

optuna_pytorch_lightning.py

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


# define the LightningModule
class LitAutoEncoder(L.LightningModule):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def training_step(self, batch, batch_idx):
        x, _ = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = nn.functional.mse_loss(x_hat, x)
        self.log("train_loss", loss)
        return loss

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


def create_model(layer_size1: int, layer_size2: int) -> LitAutoEncoder:
    encoder = nn.Sequential(nn.Linear(28 * 28, layer_size1), nn.ReLU(), nn.Linear(layer_size1, layer_size2))
    decoder = nn.Sequential(nn.Linear(layer_size2, layer_size2), nn.ReLU(), nn.Linear(layer_size2, 28 * 28))
    return LitAutoEncoder(encoder, decoder)


dataset = MNIST(os.getcwd(), download=True, transform=ToTensor())
train_loader = utils.data.DataLoader(dataset)

def objective(trial: optuna.trial.Trial):
    layer_size1 = trial.suggest_int("layer_size1", 20, 60)
    layer_size2 = trial.suggest_int("layer_size2", 3, 10)
    autoencoder = create_model(layer_size1, layer_size2)
    trainer = L.Trainer(limit_train_batches=100, max_epochs=10)
    trainer.fit(model=autoencoder, train_dataloaders=train_loader)
    trainer.logger.log_hyperparams({"layer_size1": layer_size1, "layer_size2": layer_size2})
    return trainer.callback_metrics["train_loss"].item() + (layer_size1 ** 2 + layer_size2 ** 2) ** 0.5


study = optuna.create_study()
study.optimize(objective, n_trials=100)
print(f"{study.best_params=}")

まずはモデルの定義ですが、こちらは前回と変更点ございません。
次にモデルの初期化です。前回の記事では一つのパラメータ設定だけで試していましたが、今回はOptunaの機能でパラメータサーチをするので、以下のようにして関数を通してLitAutoEndoderをインスタンスとして受け取るようにしました。

def create_model(layer_size1: int, layer_size2: int) -> LitAutoEncoder:
    encoder = nn.Sequential(nn.Linear(28 * 28, layer_size1), nn.ReLU(), nn.Linear(layer_size1, layer_size2))
    decoder = nn.Sequential(nn.Linear(layer_size2, layer_size2), nn.ReLU(), nn.Linear(layer_size2, 28 * 28))
    
    # init the autoencoder
    return LitAutoEncoder(encoder, decoder)

今回は二つのレイヤーのパラメータ数を指定するようにしたので、それらのパラメータを引数で受け取り、エンコーダとデコーダの構成を変えられるようにしています。

次にOptunaでは最適化の対象とする目的関数を定義する必要があるため、以下のように実装しました。layer_size1とlayer_size2をそれぞれ適当な範囲で取得できるようにしつつ、LightningのTrainerを実行させています。また目的関数の最適化の対象ですが、モデルのロスに加えてパラメータが大きくなることを防ぐいわゆる正則化をするために、L2ノルムを加算することで対応しています。

def objective(trial: optuna.trial.Trial):
    # We optimize the number of layers, hidden units in each layer and dropouts.
    layer_size1 = trial.suggest_int("layer_size1", 20, 60)
    layer_size2 = trial.suggest_int("layer_size2", 3, 10)
    autoencoder = create_model(layer_size1, layer_size2)
    trainer = L.Trainer(limit_train_batches=100, max_epochs=10)
    trainer.fit(model=autoencoder, train_dataloaders=train_loader)
    trainer.logger.log_hyperparams({"layer_size1": layer_size1, "layer_size2": layer_size2})
    return trainer.callback_metrics["train_loss"].item() + (layer_size1 ** 2 + layer_size2 ** 2) ** 0.5

実行してみる！

それでは先ほどのコードを実行してみましょう。

uv run optuna_pytorch_lightning.py

実行するとひたすらログが画面に表示されます。今回はトライアル回数を100に設定したので、100つのパラメータ設定の組み合わせに基づいてモデル学習を実行した結果を比較します。以下の画像はラスト3怪の組み合わせ結果の例です。結果を見ると、layer_size1が20、layer_size2が3の時に最もスコアがよくなっていることが確認できました（AutoEncoderなので、最終的な埋め込み次元が3の時いいということですね）。正則化としてL2ノルムを入れているので、どちらのパラメータも指定された範囲のうち小さいところで、できるだけパラメータ数が小さくなるようにしていることが伺えます。

まとめ

今回はOptunaとPyTorch Lightningを組み合わせて簡単なモデルのパラメータ最適化をしてみました。Optunaで簡単にパラメータ最適化をしつつ、PyTorch Lightningd簡単に学習コードを実装できるというベスト＋ベストの組み合わせを試せたと思います。皆さんもぜひ試してみてください。

それでは早速やってみる！

環境構築

ソースコードの実装

実行してみる！

まとめ

Discussion