[深層学習] SchedulerいらずのOptimizer

## 1. はじめに

皆さん、深層学習でスケジューラーを使っていて、

- どのスケジューラーを選べばいいのか分からない…
- とりあえず使ってみたけど、パラメータが多くて設定が面倒…

そんな経験はありませんか？

本記事では、**スケジューラーなしでも安定した学習が可能なOptimizer**を紹介します。

## 2. 今回使用したOptimizer

https://github.com/facebookresearch/schedule_free

schedule_freeというライブラリの RAdamScheduleFree を使用しました。

下記が元となる論文です。

https://arxiv.org/pdf/2405.15682

## 3. 実験条件

- データ: CIFAR-10(torchvisionのものを使用)
- モデル: resnetライクなモデル

<details><summary>モデルコード</summary>

```python
class BasicBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super().__init__()
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3,
                               stride=stride, padding=1, bias=False)
        self.relu = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)

        self.downsample = downsample

    def forward(self, x):
        identity = x

        out = self.bn1(x)
        out = self.relu(out)
        out = self.conv1(out)

        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        return out


class SimpleResNet(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        
        self.in_channels = 32

        self.stem = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True)
        )

        # Residual Block groups
        self.layer1 = self._make_layer(32, num_blocks=3, stride=1)  # 32x32
        self.layer2 = self._make_layer(64, num_blocks=3, stride=2)  # 16x16
        self.layer3 = self._make_layer(128, num_blocks=3, stride=2) # 8x8
        self.layer4 = self._make_layer(256, num_blocks=3, stride=2) # 4x4

        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    def _make_layer(self, out_channels, num_blocks, stride):
        downsample = None
        if stride != 1 or self.in_channels != out_channels:
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channels, out_channels,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )

        layers = []
        layers.append(BasicBlock(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels
        for _ in range(1, num_blocks):
            layers.append(BasicBlock(out_channels, out_channels))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.stem(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.pool(x)
        x = self.fc(x)
        return x
```

</details>

- 使用したデータ拡張
    - ランダム左右反転
    - RandAugment
    - Cutout

- 比較Optimizer
    - RAdam
    - RAdam + CosineAnnealing
    - RAdam(schedule free)

- データセット分割
    - train: 元学習データの80%
    - valid: 元学習データの20%
    - test: 元テストデータ

- その他ハイパーパラメータ
    - 学習率: 1e-3
    - epoch: 100
    - weight_decay: 1e-4

## 4. shcedule free optimizerを使用する際のコードの修正点

基本的に学習の最初にoptim.train()やoptim.eval()を追加するだけです。

```
# == 学習時 ==
model.train()
optim.train() # 追記
# 以下、学習コード...

# == 推論時 ==
model.eval()
optim.eval() # 追記
# 以下、推論コード
```


## 5. 実験結果

### 5.1 学習曲線

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/319102/ec1d49b2-a281-4f6e-8888-68df9207086b.png)

上記のような学習曲線が得られました。

スケジューラを使用した場合が最終的な性能は高くなりますが、schedule freeのoptimizerを使用した場合が収束は速いことが分かると思います。

### 5.2 テストデータでのaccuracyおよび混同行列

| Optimizer               | accuracy | 混同行列 |
|-------------------------|------------|------------|
| RAdam only | 0.8939 | ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/319102/3dd58051-5dcb-4b3d-9984-49cf6a14004e.png) |
| RAdam + Scheduler | 0.9323 | ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/319102/9f4e800f-1a2a-4616-9135-214635720c48.png) |
| RAdam(schedule free) | 0.9227 | ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/319102/d80210d3-4ebd-4ac8-bae3-a6440515b410.png)|

## 6. まとめ

- schedule_freeのライブラリを使用してスケジューラーを使用しないoptimizerを使用してみた
- 精度としては、optimizer + スケジューラーを使用 > schedule_free > optimizer単体 となる
- 収束の速さは、schedule_free > optimizer + スケジューラーを使用 ≒ optimizer単体 となる

## 7. 所感

今回は、scheduler_free の性能が最大にはなりませんでしたが、

- スケジューラーの設定などを調べずにとりあえずいい感じの性能を得ることが出来る
- どのスケジューラーを使用するか悩まなくてよくなる

などが魅力的だと感じました。

## 8. 意外に苦労した点

今回はCIFAR-10を使いましたが、実は「スケジューラーを使わずに学習する」以外に、
**スクラッチ実装でaccuracy 90%超えを目指す**という裏目標もありました。

しかし、既存の多くのモデルは 32×32 の画像サイズを前提としておらず、  
そのまま使うとオーバーパラメータになったり、逆に精度が出なかったりします。

そのため、小さな画像サイズに適した構造を自分で設計し、
かつ性能も出るように調整するのは地味に大変でした。

## 9. 参考文献

https://arxiv.org/pdf/2405.15682

https://github.com/facebookresearch/schedule_free

https://zenn.dev/dena/articles/6f04641801b387

https://zenn.dev/ykesamaru/articles/5e19e38990345f

## 10. 🔗 実験コード

実際に使用したコードは以下のGitHubに公開しています。  
興味のある方は、ぜひ試してみてください！

https://github.com/ryota0051/schedule_free_comparison_cifar10