PyTorch Lightning 1.1: ユースケース : 最適化 (翻訳/解説)
翻訳 : (株)クラスキャットセールスインフォメーション
作成日時 : 02/28/2021 (1.2.1)

* 本ページは、PyTorch Lightning ドキュメントの以下のページを翻訳した上で適宜、補足説明したものです：

Common Use Cases : Optimization

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

★ 無料セミナー実施中 ★ クラスキャット主催人工知能 & ビジネス Web セミナー

人工知能とビジネスをテーマにウェビナー (WEB セミナー) を定期的に開催しています。スケジュールは弊社公式 Web サイトでご確認頂けます。

お住まいの地域に関係なく Web ブラウザからご参加頂けます。事前登録 が必要ですのでご注意ください。
Windows PC のブラウザからご参加が可能です。スマートデバイスもご利用可能です。

クラスキャットは人工知能・テレワークに関する各種サービスを提供しております :

人工知能研究開発支援	人工知能研修サービス	テレワーク & オンライン授業を支援
PoC(概念実証)を失敗させないための支援 (本支援はセミナーに参加しアンケートに回答した方を対象としています。)

◆ お問合せ : 本件に関するお問い合わせ先は下記までお願いいたします。

株式会社クラスキャット セールス・マーケティング本部セールス・インフォメーション

E-Mail：sales-info@classcat.com ; WebSite: https://www.classcat.com/

Facebook: https://www.facebook.com/ClassCatJP/

ユースケース : 最適化

Lightning は最適化プロセスを管理するための 2 つのモードを提供します :

自動最適化 (AutoOpt)
手動最適化

手動最適化

強化学習、sparse coding や GAN 研究のような進んだ研究トピックのためには、最適化プロレスを手動で管理することが望ましいかもしれません。それを行なうには、以下を行ないます :

LightningModule automatic_optimization プロパティを False を返すように override します。
optimizer_idx 引数を破棄または無視します。
loss.backward() の代わりに self.manual_backward(loss) を使用します。

Note

これは究極的な柔軟性を必要とする専門家のためにだけ推奨されます。Lightning は精度と accelerator ロジックだけを処理します。ユーザは zero_grad, accumulated_grad_batches, model toggling 等が任されます。

WARNING

1.2 以前は、optimzer.step は内部的には zero_grad を呼び出していました。1.2 から、それはユーザの専門的知識に任されます。

TIP

一つの optimizer とともに accumulate_grad_batches を遂行するには、そのようなものとして行なうことができます。

TIP

self.optimizers() は LightningOptimizer オブジェクトを返します。貴方自身の optimizer に optimizer.optimizer でアクセスできます。けれども、ステップを遂行するために貴方自身の optimizer を使用するのであれば、Lightning は貴方のために accelerators と precision をサポートすることができません。

def training_step(batch, batch_idx, optimizer_idx):
    opt = self.optimizers()

    loss = self.compute_loss(batch)
    self.manual_backward(loss)
    opt.step()

    # accumulate gradient batches
    if batch_idx % 2 == 0:
        opt.zero_grad()

TIP

optimizer をモデルの forward と backward パスを遂行する closure 関数とともに提供することは良い実践です。それは殆どの optimizer についてオプションですが、closure を必要とする optimizer に切り替える場合コードを互換にします。

closure を使用した上と同じサンプルがここにあります。

def training_step(batch, batch_idx, optimizer_idx):
    opt = self.optimizers()

    def forward_and_backward():
        loss = self.compute_loss(batch)
        self.manual_backward(loss)

    opt.step(closure=forward_and_backward)

    # accumulate gradient batches
    if batch_idx % 2 == 0:
        opt.zero_grad()

# Scenario for a GAN.

def training_step(...):
    opt_gen, opt_dis = self.optimizers()

    # compute generator loss
    loss_gen = self.compute_generator_loss(...)

    # zero_grad needs to be called before backward
    opt_gen.zero_grad()
    self.manual_backward(loss_gen)
    opt_gen.step()

    # compute discriminator loss
    loss_dis = self.compute_discriminator_loss(...)

    # zero_grad needs to be called before backward
    opt_dis.zero_grad()
    self.manual_backward(loss_dis)
    opt_dis.step()

Note

LightningOptimizer は進んだユーザのために toggle_model 関数を @context_manager として提供します。それは幾つかの optimizer で勾配 accumulation を遂行したり分散設定で訓練するときに有用であり得ます。

ここにそれが何を行なうかの説明があります :

現在の optimizer を A としてそして総ての他の optimizer を B として考えます。toggling は A 専用である B からの総てのパラメータがそれらの requires_grad 属性を False に設定することを意味します。それらの元の状態はコンテキスト・マネージャを抜ける (= exit) ときにリストアされます。

勾配 accumulation を遂行するとき、accumulation フェイズの間勾配同期を遂行する必要はありません。sync_grad を False に設定すればこの同期をブロックして訓練スピードを改良します。

それをどのように使用するかのサンプルがここにあります :

# Scenario for a GAN with gradient accumulation every 2 batches and optimized for multiple gpus.

def training_step(self, batch, batch_idx, ...):
    opt_gen, opt_dis = self.optimizers()

    accumulated_grad_batches = batch_idx % 2 == 0

    # compute generator loss
    def closure_gen():
        loss_gen = self.compute_generator_loss(...)
        self.manual_backward(loss_gen)
        if accumulated_grad_batches:
            opt_gen.zero_grad()

    with opt_gen.toggle_model(sync_grad=accumulated_grad_batches):
        opt_gen.step(closure=closure_gen)

    def closure_dis():
        loss_dis = self.compute_discriminator_loss(...)
        self.manual_backward(loss_dis)
        if accumulated_grad_batches:
            opt_dis.zero_grad()

    with opt_dis.toggle_model(sync_grad=accumulated_grad_batches):
        opt_dis.step(closure=closure_dis)

自動最適化

Lightning では殆どのユーザは .backward(), .step(), .zero_grad() をいつ呼び出すかについて考えなくて構いません、何故ならば Lightning はそれを貴方のために自動化するからです。

内部的には Lightning は以下を行ないます :

for epoch in epochs:
    for batch in data:
        loss = model.training_step(batch, batch_idx, ...)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    for scheduler in schedulers:
        scheduler.step()

マルチ optimizer の場合には、Lightning は以下を行ないます :

for epoch in epochs:
  for batch in data:
     for opt in optimizers:
        disable_grads_for_other_optimizers()
        train_step(opt)
        opt.step()

  for scheduler in schedulers:
     scheduler.step()

学習率スケジューリング

使用する総ての optimizer は任意の LearningRateScheduler とペアになることができます。基本的なユースケースでは、.configure_optimizers メソッドからの 2 番目の出力としてスケジューラ (or 複数のスケジューラ) が返されるべきです :

# no LR scheduler
def configure_optimizers(self):
   return Adam(...)

# Adam + LR scheduler
def configure_optimizers(self):
   optimizer = Adam(...)
   scheduler = LambdaLR(optimizer, ...)
   return [optimizer], [scheduler]

# Two optimizers each with a scheduler
def configure_optimizers(self):
   optimizer1 = Adam(...)
   optimizer2 = SGD(...)
   scheduler1 = LambdaLR(optimizer1, ...)
   scheduler2 = LambdaLR(optimizer2, ...)
   return [optimizer1, optimizer2], [scheduler1, scheduler2]

.step() メソッドがメトリック値で条件付けられるスケジューラがあるとき (例えば ReduceLROnPlateau スケジューラ)、Lightning は configure_optimizers からの出力が辞書であるべきであることを要求します、各 optimizer に一つ、スケジューラが (その上で) 条件付けられるべきメトリックに設定されたキーワード monitor を伴います。

# The ReduceLROnPlateau scheduler requires a monitor
def configure_optimizers(self):
   return {
       'optimizer': Adam(...),
       'lr_scheduler': ReduceLROnPlateau(optimizer, ...),
       'monitor': 'metric_to_track'
   }

# In the case of two optimizers, only one using the ReduceLROnPlateau scheduler
def configure_optimizers(self):
   optimizer1 = Adam(...)
   optimizer2 = SGD(...)
   scheduler1 = ReduceLROnPlateau(optimizer1, ...)
   scheduler2 = LambdaLR(optimizer2, ...)
   return (
       {'optimizer': optimizer1, 'lr_scheduler': scheduler1, 'monitor': 'metric_to_track'},
       {'optimizer': optimizer2, 'lr_scheduler': scheduler2},
   )

Note

メトリクスは、lightning モジュールで self.log(‘metric_to_track’, metric_val) を使用してそれを単純にログ記録することにより条件付けるために利用可能になれます。

デフォルトでは、総てのスケジューラは各エポックの終了後に呼び出されます。この動作を変更するため、スケジューラ configuration は辞書として返されるべきです、これは以下のキーワードを含みます :

scheduler (required): 実際の scheduler オブジェクト
monitor (optional): 条件付けるメトリック
interval (optional): 各エポック終了後にステップするための epoch (デフォルト) か、各最適化ステップ後にステップするための step のいずれか。
frequency (optional): scheduler.step() を呼び出す間に幾つの epochs/steps をパスするべきか。デフォルトは 1 で、総ての epoch/step の後に学習率を更新することに相当します。
strict (optional): True に設定する場合、scheduler.step() を呼び出そうとする間 monitor で指定された値が利用可能であることを強制し、見つからなければ訓練を停止します。False であれば警告を与えるだけで (スケジューラを呼び出すことなく) 訓練を続けます。
name (optional): 学習率の進捗をモニタするために LearningRateMonitor コールバックを使用している場合、学習率がログ記録されるべき特定の名前として指定するためにこのキーワードが使用できます。

# Same as the above example with additional params passed to the first scheduler
# In this case the ReduceLROnPlateau will step after every 10 processed batches
def configure_optimizers(self):
   optimizers = [Adam(...), SGD(...)]
   schedulers = [
      {
         'scheduler': ReduceLROnPlateau(optimizers[0], ...),
         'monitor': 'metric_to_track',
         'interval': 'step',
         'frequency': 10,
         'strict': True,
      },
      LambdaLR(optimizers[1], ...)
   ]
   return optimizers, schedulers

複数の optimizer を利用する (like GAN)

複数の optimizer を使用するには pytorch_lightning.core.LightningModule.configure_optimizers() から > 1 optimizer を返します。

# one optimizer
def configure_optimizers(self):
   return Adam(...)

# two optimizers, no schedulers
def configure_optimizers(self):
   return Adam(...), SGD(...)

# Two optimizers, one scheduler for adam only
def configure_optimizers(self):
   return [Adam(...), SGD(...)], {'scheduler': ReduceLROnPlateau(), 'monitor': 'metric_to_track'}


Lightning は各 optimizer をシーケンシャルに呼び出します :

for epoch in epochs:
   for batch in data:
      for opt in optimizers:
         train_step(opt)
         opt.step()

   for scheduler in schedulers:
      scheduler.step()

任意のインターバルで optimizer をステップさせる

optimizer で学習率ウォームアップや変則的な (= odd) スケジューリングのようなより面白いことを行なうためには、optimizer_step() 関数を override します。

例えば、ここで 2 バッチ毎に optimizer A そして 4 バッチ毎に optimizer B をステップさせます。

Note

Trainer(enable_pl_optimizer=True) を使用するとき、.zero_grad() を呼び出す必要はありません。

def optimizer_zero_grad(self, current_epoch, batch_idx, optimizer, opt_idx):
  optimizer.zero_grad()

# Alternating schedule for optimizer steps (ie: GANs)
def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_idx, closure, on_tpu=False, using_native_amp=False, using_lbfgs=False):
    # update generator opt every 2 steps
    if optimizer_idx == 0:
        if batch_nb % 2 == 0 :
           optimizer.step(closure=closure)

    # update discriminator opt every 4 steps
    if optimizer_idx == 1:
        if batch_nb % 4 == 0 :
           optimizer.step(closure=closure)

ここでは学習率ウォームアップを追加します :

# learning rate warm-up
def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_idx, closure, on_tpu=False, using_native_amp=False, using_lbfgs=False):
    # warm up lr
    if self.trainer.global_step < 500:
        lr_scale = min(1., float(self.trainer.global_step + 1) / 500.)
        for pg in optimizer.param_groups:
            pg['lr'] = lr_scale * self.hparams.learning_rate

    # update params
    optimizer.step(closure=closure)

Note

デフォルト optimizer_step はステップを適切に遂行するために内部の LightningOptimizer に依拠します。それは TPU, AMP, accumulate_grad_batches, zero_grad 等々… を処理します。

# function hook in LightningModule
def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_idx, closure, on_tpu=False, using_native_amp=False, using_lbfgs=False):
  optimizer.step(closure=closure)

NOTE

LightningOptimizer からラップされた Optimizer にアクセスするには、以下を行ないます。

# function hook in LightningModule
def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_idx, closure, on_tpu=False, using_native_amp=False, using_lbfgs=False):

  # `optimizer is a ``LightningOptimizer`` wrapping the optimizer.
  # To access it, do as follow:
  optimizer = optimizer.optimizer

  # run step. However, it won't work on TPU, AMP, etc...
  optimizer.step(closure=closure)

最適化のために closure 関数を使用する

LBFGS のような最適化スキームを使用するとき、second_order_closure が有効にされる必要があります。デフォルトでは、この関数は次のように training_step と backward ステップをラップすることにより定義されます。

def second_order_closure(pl_module, split_batch, batch_idx, opt_idx, optimizer, hidden):
    # Model training step on a given batch
    result = pl_module.training_step(split_batch, batch_idx, opt_idx, hidden)

    # Model backward pass
    pl_module.backward(result, optimizer, opt_idx)

    # on_after_backward callback
    pl_module.on_after_backward(result.training_step_output, batch_idx, result.loss)

    return result

# This default `second_order_closure` function can be enabled by passing it directly into the `optimizer.step`
def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_idx, second_order_closure, on_tpu=False, using_native_amp=False, using_lbfgs=False):
    # update params
    optimizer.step(second_order_closure)

以上

月	火	水	木	金	土	日
1	2	3	4	5	6	7
8	9	10	11	12	13	14
15	16	17	18	19	20	21
22	23	24	25	26	27	28
29	30	31