前置き

先日、KaggleのRSNAコンペが終わり、苦戦した点やモデル学習時のポイントなどを備忘録として残しておきます。

コンペ概要

患者の乳房X線（マンモグラフィ）検査画像から乳がんかどうかを特定する二値分類の問題で、一人の患者につき、左胸にがんがあるか、右胸にがんがあるかを予測するといった内容でした。
データ数が5万件以上あり、データ不均衡対策がモデルの予測精度を上げるうえで重要でした。（がん患者の画像が全体の2%ほど）
評価指標はpF1というF1スコアの確率版です。

はじめに

Pytorch等の深層学習ライブラリは、32bit浮動小数点（FP32)を利用して計算されることが知られていますが、大規模モデルを学習する際、計算時間がかかりすぎたり、メモリの消費量が大きくなりすぎたりしてしまうという課題に直面します。
モデル精度を落とすことなく、計算速度の向上、メモリ使用量の削減が見込める混合精度という手法があります。

混合精度とは

16bitと32bitの両方の浮動小数点型を使用して、モデルの学習を高速化する手法。

<メリット>

メモリ削減により、大きなバッチサイズで学習することができる
- 学習が安定し、精度向上に繋がる

<デメリット>

Tensorコア搭載のGPUでしか使用（パフォーマンスを発揮）できない
- RTX GPU、V100、A100などが該当

最新のGPU(A100)を使用した混合精度では、通常のFP32と比べ、演算スループットが2~8倍高いと言われています。
詳細については、以下に記載されています。

混合精度の実装

今回、コンペで実際に使用したコードを参考に説明します。

必要なライブラリのimport

torch.cuda.ampモジュールを使用します。

import torch
from torch.cuda.amp import GradScaler, autocast

スケーラーの定義

勾配情報をスケールするためのスケーラーを定義します。

scaler = GradScaler()

バッチ毎に学習

以下が混合精度を組み込んだ学習時のコードになります。
学習時のみ勾配計算を行い、検証時には勾配の更新を実施しないことに注意が必要です。

def run_train(
    model,
    train_dataloader,
    optimizer,
    scheduler,
    cfg,
    scaler,
    writer,
    epoch,
    iteration,
    step,
    loss_function,
):
    model.train()
    losses = []
    progress_bar = tqdm(range(len(train_dataloader)))
    tr_it = iter(train_dataloader)
    print(tr_it)

    all_outputs, all_labels = [], []

    for itr in progress_bar:
        batch = next(tr_it)
        inputs, labels = batch["image"].to(cfg.device), batch["label"].float().to(cfg.device)
        iteration += 1

        step += cfg.batch_size
        # 学習時のみ勾配計算を行う
        torch.set_grad_enabled(True)
        with autocast():
            outputs = model(inputs)
            loss = loss_function(outputs, labels)
        losses.append(loss.item())

        outputs = list(torch.sigmoid(outputs).detach().cpu().numpy())
        labels = list(labels.detach().cpu().numpy())

        all_outputs.extend(outputs)
        all_labels.extend(labels)

        # スケールした勾配を作る
        scaler.scale(loss).backward()
        # 勾配爆発を防ぐために勾配をクリップする
        torch.nn.utils.clip_grad_norm_(model.parameters(), 10.0)
        # optimizerに割り当てられた勾配をアンスケールしてパラメータの更新
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()
        scheduler.step()

        progress_bar.set_description(f"loss: {np.mean(losses):.2f} lr: {scheduler.get_last_lr()[0]:.6f}")
    score = pfbeta(all_labels, all_outputs, 1.0)[0]
    auc = roc_auc_score(all_labels, all_outputs)
    print("Train F1: ", score, "AUC: ", auc)

開発環境

今回は、Google Colaboratory pro+ A100（プレミアムGPU）を使用して学習しましたが、Google Colaboratoryの仕様変更により、今後Kaggleに取り組む上でGoogle Colaboratoryでの開発は厳しいと感じました。

Google Colaboratory 仕様変更内容

月額定額でGPUの無制限使用　-> 月額一定のコンピューティングユニットを購入

Google Colaboratoryの代替として、Paperspace Gradientというサービスが個人的に良さそうに感じたため、今後はこちらを使用していきたいと思います。

Paperspace Gradientの特徴

Free, Pro, Growthの3つのプランがある
月額定額でハイエンドGPUを無償で使用することができる
- 但し、最大6時間のインスタンス実行制限がある

最後に

最後まで読んでいただきありがとうございました。
今後も低コストで混合精度を利用できる開発環境を追求していきたいと思います。

参考文献

https://www.tensorflow.org/guide/mixed_precision?hl=ja

Pytorchで混合精度を使用してみる

前置き