小ネタ：Pytorch で Automatic Mixed Precision (AMP) の ON/OFF をするときの話

最近気付いたのでメモ。長くなってしまったので、結論だけ見たい場合はまとめまで読み飛ばしてください。

まえおき

NN を学習する際の GPU メモリの使用量軽減や学習速度の向上手段として混合精度学習(Mixed Precision Training) は欠かせません。pytorch では torch.cuda.amp モジュールを用いることでとてもお手軽に使うことが可能です。

以下は official docs に Typical Mixed Precision Training と題して載っている例ですが、 model の forward と loss の計算を amp.autocast の with 文中で行い、loss の backward と optimizer の step に amp.GradScaler を介在させています*1。

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Creates a GradScaler once at the beginning of training.
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # Runs the forward pass with autocasting.
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
        # Backward passes under autocast are not recommended.
        # Backward ops run in the same dtype autocast chose for corresponding forward ops.
        scaler.scale(loss).backward()

        # scaler.step() first unscales the gradients of the optimizer's assigned params.
        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,
        # otherwise, optimizer.step() is skipped.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        scaler.update()

因みに AMP を使わない場合は以下のようになると思いますが、引用したものと比較すると(コメントの部分を除いて)変更はごくごく一部であることがわかります。

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        output = model(input)
        loss = loss_fn(output, target)

        loss.backward()
        optimizer.step()

このちょっとの変更でGPU消費量が半分くらいになり計算速度も2倍ぐらいになるので、使わない理由は特に無いですよね*2。

本題

混合精度学習を使わない理由は無いと言いましたが、対応していないGPUを使う場合や計算精度が非常に重要な場合といった使わない(使えない)場面も存在します。

そういった場合にいちいち code を書き換えるのも面倒なので、上の例になぞらえると僕の場合は以下のような感じで分岐を書いていました。ここで、use_amp は混合精度学習を行う際に True にする変数です。

model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
scaler = GradScaler() if use_amp else None

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        if use_amp: # 混合精度学習をする
            with autocast():
                output = model(input)
                loss = loss_fn(output, target)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:       # 通常の学習をする
            output = model(input)
            loss = loss_fn(output, target)
            loss.backward()
            optimizer.step()

冒頭に挙げた例の違う部分をそのまま分岐で書いてるわけですが、うーんなんというかダサいですね。全然違うことを書いてるんならいいんですけど code としてはほぼいっしょなんだよなあ。何とかならないんでしょうか？

解決策(?)

何と無しに torch.cuda.amp.autocast の docs を眺めていたら気になる記述がありました。

autocast(enabled=False) subregions can be nested in autocast-enabled regions. Locally disabling autocast can be useful, for example, if you want to force a subregion to run in a particular dtype. ...

おまじない的に使っていたせいで知らなかったのですが autocast() には引数 enabled があり default は True になっています。上の記述は with autocast(enabled=True) で autocast が有効になっている範囲の中でも with autocast(enabled=False) でくくった範囲は autocast を無効化出来るという話なのですが、「なら autocast が有効でない状態での with autocast(enabled=False) は何もしないのでは？」という発想に至ります。

実装(以下)を見る限りでも、with 文の中に入るときに元の状態(self.prev)を覚えておいて(self._enabled が True なら) autocast を有効にし、with 文から出るときに元の状態に戻しているのでそれっぽいです。

    def __enter__(self):
        self.prev = torch.is_autocast_enabled()
        torch.set_autocast_enabled(self._enabled)
        torch.autocast_increment_nesting()

    def __exit__(self, *args):
        # Drop the cache when we exit to a nesting level that's outside any instance of autocast.
        if torch.autocast_decrement_nesting() == 0:
            torch.clear_autocast_cache()
        torch.set_autocast_enabled(self.prev)
        return False

torch.autocast_increment_nesting が少し気になるところですが、まあ autocast が無効化されてるときは何もしないと思いたい。

というわけで、model の forward と loss の計算に関しては分岐せずとも with autocast(enabled=use_amp) としてしまえばよしなにしてくれそうです。あとは GradScaler をどうにかすればいいのですが、実はこっちに関してはもっと分かりやすいです。GradScaler はインスタンス化する瞬間に enabled という引数を渡すことが出来て、この引数の値でメソッドの挙動が変化します。

scale

Returns scaled outputs. If this instance of GradScaler is not enabled, outputs are returned unmodified.

要するに scaler.scale(loss).backward() は enabled=False の際にはただの loss.backward() です。
step

実装を見るとわかりますが、enabled=False の時には scaler.step(optimizer) は単純に optimizer.step() をします。
update

同じく実装を見るとわかりますが、enabled=False の時には何もしません。

したがってscaler = GradScaler(enabled=use_amp) としてしまえば AMP を使わない場合でも何も影響しないことがわかりました。

まとめ

以上のことから僕の書いていたダサい分岐は消滅しました。やったぜ。

model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
scaler = GradScaler(enabled=use_amp)   # 初期化時にAMPを使うかを渡す

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        with autocast(is_enabled=use_amp):  # 呼び出すときにAMPを使うかを渡す
            output = model(input)
            loss = loss_fn(output, target)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

AMP を使うときの書き方を知ってから何も考えずおまじない的に autocast と GradScaler を書いていたのですがもっと早く docs を読むべきでした。自分の使うもののドキュメントや実装を確かめて動き方を把握するのはやっぱり大事ですね。

autocast についてはちょっと自信ないので、間違っている場合はコメントして頂ければと思います。ではでは。

*1:おそらく簡単のために input と target を gpu device に送る部分などが省略されているので注意

*2:計算精度が必要になってくる場合は別

俵言

しがない社会人が書く、勉強とかのこと。最近は機械学習や kaggle 関連がメイン。

小ネタ：Pytorch で Automatic Mixed Precision (AMP) の ON/OFF をするときの話

まえおき

本題

解決策(?)

まとめ