PytorchでRAdamを使う -最適化アルゴリズムまとめ-

こんにちは。

現役エンジニアの”はやぶさ”@Cpp_Learningです。最近は Pytorch を使って深層学習を楽しんでいます。

今回は以下の内容について勉強したので、備忘録も兼ねて本記事を書きます。

本記事に含まれる内容

最適化アルゴリズム（Optimizer）の情報整理
Pytorch（Pytorch Lightning）から様々なOptimizerを使う方法
Pytorch LightningでRAdamを使う

Contents

1 最適化アルゴリズム（Optimizer）とは
2 Pytorchで様々な最適化アルゴリズム（Optimizer）を使う
- 2.1 pytorch-optimizerの使い方
3 実践！Pytorch LightningでRAdamを使う
4 ベストプラクティスを求めて
5 まとめ

最適化アルゴリズム（Optimizer）とは

最適化アルゴリズム？なにそれ美味しいの？

という人は、以下の書籍でDeep Learningについて勉強すると良いですよ(*･ω･)ﾉ♪

ゼロから作るDeep Learning ―Pythonで学ぶディープラーニングの理論と実装

created by Rinker

以下のブログ記事で勉強するのも良いと思います。

最適化アルゴリズム（Optimizer）選定

どの最適化アルゴリズムを使うべきか分からない

という人は、各手法を比較している以下の記事が参考になると思います。

SGD Momentumは良いぞ～

個人的には、最適化アルゴリズムに悩むときは SGD Momentum 使っておけば良いと考えています（以下の資料がとても勉強になります）

SGD Momentum の気持ちを理解する

RAdamとは

ただ最近、RAdamが良いぞ～という情報をキャッチしたので気になっています。RAdamについては論文や以下の解説記事を読むと良いです。

ついにAdamを超えた！最新の最適化アルゴリズム「RAdam」解説｜Qiita

Pytorchで様々な最適化アルゴリズム（Optimizer）を使う

Pytorch公式は様々な最適化アルゴリズム（Optimizer）をサポートしていますが、その中にRAdam はありません（2020/03/08時点）

そのため、RAdamを試す場合は自作する必要があります。論文の筆者が公開しているソースコードを使っても良いのですが…

RAdam含め多数の最適化アルゴリズムを実装しているリポジトリを見つけました（下記GitHub参照）。

pytorch-optimizer｜GitHub

しかもシンプルで使い易い！ありがとうございます。使わせていただきます。

pytorch-optimizerの使い方

README に丁寧な説明があるし、examples（サンプルコード）も用意してあるので、pytorchでoptimizerを使ったことのある人なら、直ぐに試せると思います。ただ、個人的に…

はやぶさ

学習用のループ処理を書きたくない…スッキリとした学習コードを書きたい…

と思ったので、pytorch-lightningを使うことにしました（pytorch-lightningの基本的な使い方は、以下の記事にまとめておきました）。

PyTorch Lightning入門から実践まで

PyTorch Lightning入門から実践まで -自前データセットで学習し画像分類モデルを生成-ディープラーニングフレームワークPytorchの軽量ラッパー”pytorch-lightning”の入門から実践までのチュートリアル記事を書きました。自前データセットを学習して画像分類モデルを生成し、そのモデルを使って推論するところまでソースコード付で解説しています。...

実践！Pytorch LightningでRAdamを使う

Pytorch LightningからRAdamを使ってみます。

Optimizer処理については ピュアPytorch と Pytorch Lightning どちらも同じコードを使えます（なので、Pytorch Lightnigを使わない人も参考になるかと）

インストール

使用するライブラリ一覧は以下の通りです。

Requirements

torch 1.4.0
torchvision 0.5.0
torch-optimizer 0.0.1a9
pytorch-lightning 0.7.1

※Google Colaboratoryで動作確認しました（2020/03/08）

Google Colaboratoryなら、以下のコマンドでインストールするだけで環境構築完了です。

pip install torch_optimizer pytorch-lightning

以降からソースコードを書いていきます。

Import

最初はimportから

import os
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
# import torch.optim as optim
import torch_optimizer as optim

from torchvision import datasets
from torchvision import transforms
from torchvision import utils

import pytorch_lightning as pl

import os

from dataclasses import dataclass

import torch

import torch.nn as nn

import torch.nn.functional as F

import torch.utils.data

# import torch.optim as optim

import torch_optimizer as optim

from torchvision import datasets

from torchvision import transforms

from torchvision import utils

import pytorch_lightning as pl

CNN設計

Pytorchの公式チュートリアルを参考にCNNを設計します。

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 4 * 4, 115)
        self.fc2 = nn.Linear(115, 84)
        self.fc3 = nn.Linear(84, 10)
 
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        output = F.log_softmax(x, dim=1)
        return output

class Net(nn.Module):

def __init__(self):

super().__init__()

self.conv1 = nn.Conv2d(1, 6, 5)

self.pool = nn.MaxPool2d(2, 2)

self.conv2 = nn.Conv2d(6, 16, 5)

self.fc1 = nn.Linear(16 * 4 * 4, 115)

self.fc2 = nn.Linear(115, 84)

self.fc3 = nn.Linear(84, 10)

def forward(self, x):

x = self.pool(F.relu(self.conv1(x)))

x = self.pool(F.relu(self.conv2(x)))

x = x.view(-1, 16 * 4 * 4)

x = F.relu(self.fc1(x))

x = F.relu(self.fc2(x))

x = self.fc3(x)

output = F.log_softmax(x, dim=1)

return output

Config(パラメータ設定)

最近、Data Classesの存在に気づいたので、config（パラメータ設定）に使ってみます。

@dataclass
class Config:
    batch_size: int = 128
    epochs: int = 10
    lr: float = 1e-3
    momentum: float = 0.9
    gamma: float = 0.7
    betas: float = (0.9, 0.999)
    eps: float = 1e-8
    weight_decay: float = 0

@dataclass

class Config:

batch_size: int = 128

epochs: int = 10

lr: float = 1e-3

momentum: float = 0.9

gamma: float = 0.7

betas: float = (0.9, 0.999)

eps: float = 1e-8

weight_decay: float = 0

※本記事では未使用のパラメータも定義しています

Data Classesの使い方については、以下の記事が参考になります。

System設計

PyTorch Lightningでは、以下のような学習システムを設計します。

class CoolSystem(pl.LightningModule):

    def __init__(self, conf):
        super(CoolSystem, self).__init__()
        self.model = Net()
        self.dir = os.getcwd()

    def forward(self, x):
        x = self.model(x)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = F.cross_entropy(y_hat, y)
        tensorboard_logs = {'train_loss': loss}
        return {'loss': loss, 'log': tensorboard_logs}

    def validation_step(self, batch, batch_idx):
        x, y = batch
        out = self.forward(x)
        return {'val_loss': F.cross_entropy(out, y)}

    def validation_end(self, outputs):
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        tensorboard_logs = {'val_loss': avg_loss}
        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}

    def configure_optimizers(self):
        # optimizer = torch.optim.SGD(self.parameters(), lr=conf.lr, momentum=conf.momentum)
        # optimizer = torch.optim.Adam(self.parameters(), lr=conf.lr)
        optimizer = optim.RAdam(
            self.parameters(),
            lr=conf.lr,
            betas=conf.betas,
            eps=conf.eps,
            weight_decay=conf.weight_decay,
        )
        return optimizer

    @pl.data_loader
    def train_dataloader(self):
        return torch.utils.data.DataLoader(
            datasets.MNIST(self.dir, train=True, download=True, transform=transforms.ToTensor()),
            batch_size=conf.batch_size, shuffle=True)

    @pl.data_loader
    def val_dataloader(self):
        return torch.utils.data.DataLoader(
            datasets.MNIST(self.dir, train=False, download=True, transform=transforms.ToTensor()),
            batch_size=conf.batch_size, shuffle=False)

class CoolSystem(pl.LightningModule):

def __init__(self, conf):

super(CoolSystem, self).__init__()

self.model = Net()

self.dir = os.getcwd()

def forward(self, x):

x = self.model(x)

return x

def training_step(self, batch, batch_idx):

x, y = batch

y_hat = self.forward(x)

loss = F.cross_entropy(y_hat, y)

tensorboard_logs = {'train_loss': loss}

return {'loss': loss, 'log': tensorboard_logs}

def validation_step(self, batch, batch_idx):

x, y = batch

out = self.forward(x)

return {'val_loss': F.cross_entropy(out, y)}

def validation_end(self, outputs):

avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()

tensorboard_logs = {'val_loss': avg_loss}

return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}

def configure_optimizers(self):

# optimizer = torch.optim.SGD(self.parameters(), lr=conf.lr, momentum=conf.momentum)

# optimizer = torch.optim.Adam(self.parameters(), lr=conf.lr)

optimizer = optim.RAdam(

self.parameters(),

lr=conf.lr,

betas=conf.betas,

eps=conf.eps,

weight_decay=conf.weight_decay,

)

return optimizer

@pl.data_loader

def train_dataloader(self):

return torch.utils.data.DataLoader(

datasets.MNIST(self.dir, train=True, download=True, transform=transforms.ToTensor()),

batch_size=conf.batch_size, shuffle=True)

@pl.data_loader

def val_dataloader(self):

return torch.utils.data.DataLoader(

datasets.MNIST(self.dir, train=False, download=True, transform=transforms.ToTensor()),

batch_size=conf.batch_size, shuffle=False)

ポイントは以下の通りです。

パラメータを定義した conf があること以外は、Pytorch Lightningの基本的なコードと同じ
Optimizerのコードは configure_optimizers() に書きます（ピュアPyTorchと同じコードを書けばOKです）
あとからSGD MomentumやAdamと比較するため、コードを書いておく（30~31行目参照）

学習

以下のステップで学習します。

Config()のインスタンスを生成
ハイパーパラメータチューニング
CoolSystemのインスタンスを生成
Trainerを定義
fitで学習

今回は以下のようにしました。

conf = Config()
cool_model = CoolSystem(conf)

trainer = pl.Trainer(
                        max_nb_epochs = conf.epochs,
                        gpus = 0 if torch.cuda.is_available() else None
                     )

trainer.fit(cool_model)