PyTorch LambdaLRで学習率を賢く調整!詳細解説と代替方法


このチュートリアルでは、torch.optim.LambdaLR の仕組みと使用方法を以下の通り詳細に解説します。

torch.optim.LambdaLR は、PyTorch Optimizerオブジェクトを受け取り、各エポックにおける学習率を決定するための関数(ラムダ式)を指定することで、学習率を制御します。このラムダ式は、現在のエポック番号を引数として受け取り、0から1の間の値を返すように設計されています。この値は、初期学習率と掛け合わされることで、エポックごとに変化する学習率が決定されます。

torch.optim.LambdaLR の構成要素

torch.optim.LambdaLR を構成する要素は以下の通りです。

  • last_epoch
    過去のエポック数。デフォルトは -1 であり、この場合、現在のエポックが最初のエポックとして扱われます。
  • lr_lambda
    学習率を更新するための関数(ラムダ式)または関数リスト。Optimizer における各パラメータグループに対して個別に設定することも可能です。
  • optimizer
    学習率を調整したい対象となる Optimizerオブジェクト

torch.optim.LambdaLR の使用方法

torch.optim.LambdaLR の基本的な使用方法は以下の通りです。

from torch.optim.lr_scheduler import LambdaLR

def lr_lambda(current_epoch):
    # エポックごとに学習率を調整するラムダ式を定義
    if current_epoch < 10:
        return 1.0
    elif current_epoch < 20:
        return 0.5
    else:
        return 0.1

optimizer = torch.optim.Adam(model.parameters())
scheduler = LambdaLR(optimizer, lr_lambda)

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

    scheduler.step()  # 各エポック後に学習率を更新

上記の例では、最初の10エポックは学習率1.0で学習し、その後10エポックは学習率0.5、最後の10エポックは学習率0.1で学習するというスケジュールになっています。

torch.optim.LambdaLR の活用例

torch.optim.LambdaLR は、以下のような様々な学習率調整シナリオに活用できます。

  • 余弦類似スケジュール
    学習率を余弦類似関数で変化させて、学習初期と後半で学習率を大きく変化させることができます。
  • コサインアニーリング
    学習率をコサイン関数で滑らかに変化させて、より安定した収束を目指すことができます。
  • 学習後半における学習率の減衰
    モデルが過学習しやすくなるのを防ぐために、学習後半の学習率を徐々に減衰させることができます。
  • 初期エポックにおける学習率の急激な増加
    モデルが初期段階で学習しやすくなるように、最初の数エポックの学習率を大きく設定することができます。
  • 学習率の調整は、モデルのパフォーマンスに大きな影響を与える可能性があります。適切な学習率スケジュールを見つけるためには、様々な設定を試して検証することが重要です。
  • torch.optim.LambdaLR は、学習率調整のためのシンプルなツールであり、より高度な調整には他のスケジューラの方が適している場合があります。


import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import LambdaLR

# モデル定義
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)

# データ準備
x = torch.randn(100, 10)
y = torch.randn(100, 1)

# モデルとオプティマイザの生成
model = MyModel()
optimizer = optim.Adam(model.parameters())

# 学習率スケジュール定義
def lr_lambda(current_epoch):
    # エポックごとに学習率を調整するラムダ式を定義
    if current_epoch < 10:
        return 1.0
    elif current_epoch < 20:
        return 0.5
    else:
        return 0.1

# 学習率スケジューラ生成
scheduler = LambdaLR(optimizer, lr_lambda)

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

    # 各エポック後に学習率を更新
    scheduler.step()

# モデルの保存
torch.save(model.state_dict(), 'my_model.pth')
  • scheduler.step() は、各エポックの終わりに呼び出され、学習率を更新します。
  • LambdaLR オブジェクトは、optimizerlr_lambda 関数を使用して作成されています。
  • lr_lambda 関数は、エポック番号に応じて学習率を調整するラムダ式として定義されています。


StepLR


  • 特徴
    • 特定のステップ間隔で学習率を一定の値で減衰させるシンプルな方法です。
    • コードが簡潔で分かりやすい。
from torch.optim.lr_scheduler import StepLR

optimizer = optim.Adam(model.parameters())
scheduler = StepLR(optimizer, step_size=20, gamma=0.5)  # 20エポックごとに学習率を0.5倍に減衰

MultiStepLR


  • 特徴
    • 指定されたエポックのリストで学習率を個別に設定することができます。
    • より柔軟な学習率調整が可能。
from torch.optim.lr_scheduler import MultiStepLR

optimizer = optim.Adam(model.parameters())
milestones = [20, 50, 80]  # 指定されたエポック
gammas = [0.5, 0.25, 0.1]  # 各エポックにおける学習率の減衰率
scheduler = MultiStepLR(optimizer, milestones=milestones, gammas=gammas)

CosineAnnealingLR


  • 特徴
    • 学習率を滑らかにコサイン関数で減衰させ、より安定した収束を目指すことができます。
    • 過学習を防ぐ効果が期待できます。
from torch.optim.lr_scheduler import CosineAnnealingLR

optimizer = optim.Adam(model.parameters())
scheduler = CosineAnnealingLR(optimizer, T_max=100)  # 100エポックかけて学習率を0に減衰

ReduceLROnPlateau


  • 特徴
    • 検証損失が一定回以上悪化しない限り学習率を維持し、悪化した場合は学習率を減衰させる動的なスケジューリングです。
    • 過学習を防ぎ、汎化性能を向上させる効果が期待できます。
from torch.optim.lr_scheduler import ReduceLROnPlateau

optimizer = optim.Adam(model.parameters())
scheduler = ReduceLROnPlateau(optimizer, mode='min', patience=10, factor=0.5)  # 検証損失が10エポック連続で悪化しない限り学習率を維持し、悪化した場合は0.5倍に減衰

カスタムスケジューラ


  • 特徴
    • 上記のスケジューラよりも柔軟性に優れ、独自の学習率調整ロジックを記述することができます。
    • 高度な学習率調整が必要な場合に有効です。
class MyScheduler(object):
    def __init__(self, optimizer, **kwargs):
        # スケジューラの初期化処理
        pass

    def step(self, epoch):
        # エポックに基づいて学習率を更新する処理
        pass

optimizer = optim.Adam(model.parameters())
scheduler = MyScheduler(optimizer)

上記以外にも、以下のような学習率調整スケジューラが PyTorch に用意されています:

  • OneCycleLR
  • CyclicLR
  • ExponentialLR

それぞれのスケジューラは異なる特性を持っているので、学習タスクやモデルに合わせて最適なものを選択することが重要です。

  • 実装の容易さ
    • StepLRMultiStepLR はコードが簡潔で分かりやすいです。
    • CosineAnnealingLRReduceLROnPlateau は若干複雑ですが、デフォルトのパラメータで良い結果を得られることが多いです。
    • カスタムスケジューラは実装に時間がかかりますが
  • 必要な柔軟性
    • 単純な減衰スケジュールであれば StepLRMultiStepLR が適しています。
    • より滑らかな減衰や動的な調整が必要であれば CosineAnnealingLRReduceLROnPlateau が適しています。
    • 高度な調整が必要であればカスタムスケジューラを検討しましょう。