PyTorchで検証指標に基づいて学習率を自動調整する: ReduceLROnPlateauの使い方とサンプルコード


torch.optim.ReduceLROnPlateau は、PyTorchにおける学習率調整のためのツールです。検証指標の改善が一定期間停止した場合に、学習率を自動的に減衰させる機能を提供します。モデルの学習が停滞していることを検知し、学習率を調整することで、最適解に効率的に到達することを目的としています。

仕組み

ReduceLROnPlateau は、以下の情報を元に学習率を調整します。

  • "verbose"
    学習率調整に関する情報を表示するかどうかを指定します。
  • "factor"
    学習率を減衰させる際の係数。デフォルトは0.1です。
  • "patience"
    指標の改善が停止したとみなすまでのエポック数。
  • "mode"
    検証指標の改善をどのように判断するかを指定します。"min" (デフォルト) と "max" の2つのモードがあり、それぞれ指標の最小値と最大値の更新に基づいて判断します。
  • 検証指標
    モデルの性能を評価するための指標。損失関数や精度指標などが一般的です。

具体的な使い方

以下のコード例は、ReduceLROnPlateau を使って学習率を調整する例です。

import torch
from torch.optim import SGD, ReduceLROnPlateau

# モデルと最適化アルゴリズムの定義
model = ...
optimizer = SGD(model.parameters(), lr=0.1)

# 検証指標の定義
criterion = ...

# ReduceLROnPlateauの作成
scheduler = ReduceLROnPlateau(optimizer, 'min', patience=10, factor=0.1, verbose=True)

# トレーニングループ
for epoch in range(num_epochs):
    # ... (トレーニング処理)

    # 検証
    with torch.no_grad():
        val_loss = ...

    # 検証指標の更新
    scheduler.step(val_loss)

このコードでは、以下の処理が行われます。

  1. モデルと最適化アルゴリズムを定義します。
  2. 検証指標を定義します。
  3. ReduceLROnPlateau を作成します。
  4. トレーニングループ内で、モデルのトレーニングと検証を行います。
  5. 検証指標に基づいて学習率を調整します。

ReduceLROnPlateau の利点

  • モデルの収束を促進できる
    学習率が小さすぎると、モデルの学習が遅くなります。ReduceLROnPlateau は、学習率を適切な値に保つことで、モデルの収束を促進することができます。
  • 過学習を防ぐことができる
    学習率が大きすぎると、モデルが訓練データに過剰に適合し、汎化性能が低下する可能性があります。ReduceLROnPlateau は、学習率を適切な値に保つことで、過学習を防ぐことができます。
  • 学習率の調整を自動化できる
    手動で学習率を調整する必要がなくなり、効率的にモデルを学習することができます。
  • 学習率の減衰率
    "factor" パラメータで学習率の減衰率を指定しますが、小さすぎると学習率の調整効果が小さくなり、大きすぎると学習率が急激に減衰してしまいます。
  • "patience" パラメータの設定
    "patience" パラメータを小さすぎると、学習率が早すぎるタイミングで減衰されてしまう可能性があります。逆に、大きすぎると、学習率が減衰されないまま、モデルが学習停止に陥る可能性があります。
  • 検証指標の選択
    適切な検証指標を選択しないと、ReduceLROnPlateau が正しく動作しない可能性があります。


import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# デバイスの設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# データセットの読み込み
train_dataset = datasets.MNIST(root="data", train=True, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

test_dataset = datasets.MNIST(root="data", train=False, download=True, transform=transforms.ToTensor())
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# モデルの定義
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(960, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = x.view(-1, 960)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        return x

model = Model().to(device)

# 損失関数と最適化アルゴリズムの定義
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# ReduceLROnPlateauの作成
scheduler = ReduceLROnPlateau(optimizer, 'min', patience=10, factor=0.1, verbose=True)

# トレーニングループ
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(train_loader, 1):
        images, labels = data[0].to(device), data[1].to(device)

        # 勾配をゼロ化
        optimizer.zero_grad()

        # 順伝播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 逆伝播
        loss.backward()

        # 重み更新
        optimizer.step()

        # 損失の記録
        running_loss += loss.item()

        if i % 100 == 0:
            print(f"[Epoch {epoch + 1}, {i * 64}/{len(train_dataset)}] Loss: {running_loss / 100:.4f}")
            running_loss = 0.0

    # 検証
    test_loss = 0.0
    correct = 0
    with torch.no_grad():
        for data in test_loader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = model(images)
            test_loss += criterion(outputs, labels).item()
            _, predicted = torch.max(outputs.data, 1)
            correct += (predicted == labels).sum().item()

    test_loss /= len(test_loader)
    test_acc = correct / len(test_loader)
    print(f"Test Epoch {epoch + 1}: Avg. loss: {test_loss:.4f}, Accuracy: {test_acc:.4f}")

    # 検証指標に基づいて学習率を調整
    scheduler.step(test_loss)
  1. MNIST データセットを訓練データと検証データに分割します。
  2. モデルを定義します。
  3. 損失関数と最適化アルゴリズムを定義します。
  4. ReduceLROnPlateau を作成します。
  5. トレーニングループ内で、モデルのトレーニングと検証を行います。


CosineAnnealingLR

CosineAnnealingLR は、学習率を余弦関数で減衰させるスケジューラです。 最初は高い学習率で学習を開始し、徐々に0に向かって減衰していきます。 この方法は、初期段階でモデルを迅速に学習させ、その後は微調整を行うのに適しています。

from torch.optim.lr_scheduler import CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs)

MultiStepLR

MultiStepLR は、指定されたエポックで学習率をステップ状に減衰させるスケジューラです。 複数の減衰率を定義することができ、学習状況に合わせて柔軟に調整することができます。

from torch.optim.lr_scheduler import MultiStepLR

milestones = [30, 60, 90]
scheduler = MultiStepLR(optimizer, milestones=[0.1, 0.05, 0.01])

LambdaLR

LambdaLR は、任意の関数に基づいて学習率を更新するスケジューラです。 カスタムな減衰スケジュールを実装したい場合に役立ちます。

from torch.optim.lr_scheduler import LambdaLR

def lr_lambda(epoch):
    return 0.95 ** epoch

scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)

EarlyStopping

EarlyStopping は、検証指標の改善が一定期間停止した場合にトレーニングを終了させる手法です。 過学習を防ぎ、モデルの汎化性能を向上させるのに役立ちます。

from torch.utils.data import EarlyStopping

early_stopping = EarlyStopping(patience=10, verbose=True)

for epoch in range(num_epochs):
    # ... (トレーニング処理)

    # 検証
    with torch.no_grad():
        val_loss = ...

    early_stopping(val_loss)
    if early_stopping.early_stop:
        break

カスタムスケジューラ

上記以外にも、torch.optim.lr_scheduler クラスを継承して独自のスケジューラを作成することもできます。 より複雑な減衰スケジュールを実装したい場合や、他の学習率調整手法と組み合わせたい場合などに有効です。

選択の指針

どの代替方法が適切かは、データセット、モデル、タスクなどの状況によって異なります。 以下に、いくつかの指針を紹介します。

  • 過学習の防止
    EarlyStopping は、過学習を防ぎたい場合に適しています。
  • 柔軟な減衰スケジュール
    LambdaLR またはカスタムスケジューラは、より複雑な減衰スケジュールを実装したい場合に適しています。
  • シンプルな減衰スケジュール
    CosineAnnealingLR または MultiStepLR は、シンプルな減衰スケジュールで学習率を調整したい場合に適しています。