PyTorchを使用した自動的な混合精度

PyTorch

紹介

大きな深層学习モデルはより多くの計算能力とメモリリソースを必要とします。新しい技術の開発により、深層 neural networkのより早い学習が実現しています。FP32 (完全精度の浮動小数点数形式)を使用する代わりに、FP16 (半精度の浮動小数点数形式)を使用することができます。研究者たちは、これらを併用する方がより良い結果を得ることを発見しました。

半精度の学習を行っていながらも、单精度のネットワークの精度の多くを保持することができる混合精度を実現します。混合精度技術とは、この方法が单精度と半精度の表現の両方を利用することを指します。

このPyTorchでのAutomatic Mixed Precision (Amp) トレーニングの概要で、技術の機能を示し、Ampを使用する手順を詳細に説明し、Amp技術のより高度な应用について讨議します。それには、用户が自分のコードに後で統合するためのコード骨格が含まれます。

前提条件

PyTorchの基本知識: PyTorchに熟悉しており、テンサー、モジュール、トレーニングループなどの主要な概念を把握していること。

深層学習の基礎理解: ニューラルネットワーク、反伝播、最適化などの概念。

混合精度トレーニングの知識: 混合精度トレーニングの利点と欠点を理解することが重要であり、メモリ使用量を減らし、計算を速めることが含まれます。

互換性のあるハードウェアへのアクセス: 混合精度をサポートしているGPUを持っていること。例えば、Tensor Coresを持つNVIDIA GPU（たとえば、Volta、Turing、Ampereアーキテクチャ）です。

PythonとCUDAのセットアップ: PyTorchをインストールし、GPU加速のためにCUDAを設定された機能しているPython環境です。

混合精度の概要

Deep learningのほとんどのフレームワークと同様に、PyTorchは通常32位浮動小数点データ（FP32）でトレーニングします。しかし、FP32は常に成功に必要ではありません。FP32が時間とメモリを多く消費する少数の操作には16位浮動小数点データを使用することができます。

その結果、NVIDIAのエンジニアらは、少ない操作にFP32で混合精度トレーニングを行うことができ、ネットワークのほとんどがFP16で走ることができる技術を開発しました。

可能な限りfloat16データ型を使用するモデル変換を行います。
float32の親重みを保持して、それぞれのイテレーションで重み更新を集めます。
損失スケーリングを使用して、微小な勾配値を保持します。

混成精度のトレーニングには、PyTorchには既に搭載されている丰富な機能があります。
モジュールのパラメーターは.half()メソッドを呼び出すとFP16に変換され、テンサンスのデータは.half()メソッドを呼び出すとFP16に変換されます。これらのモジュールやテンサンスに対する操作は、速いFP16算術を使用して実行されます。NVIDIAの数学ライブラリ（cuBLASおよびcuDNN）はPyTorchと密接に互換性があります。FP16パイプラインからのデータは、Tensor Coresを使用してGEMMやコンボルューションを行います。cuBLASでTensor Coresを使用するには、GEMMの次元([M, K] x [K, N] -> [M, N])は8の倍数である必要があります。

Apexの紹介

Apexの混成精度ユーティリティは、トレーニング速度を上昇させるだけでなく、单精度トレーニングの精度と安定性を保持するために設計されています。ApexはFP16またはFP32で操作を行い、マスターパラメーターの変換を自動的に行い、損失の缩放を自動的に行います。

Apexは研究者がミックス精度の学習をモデルに組み込むことを簡素化するために作られました。AmpはApexの機能の1つ、自動マイクス精度の略称で、軽量のPyTorchの拡張です。ユーザーはAmpを使用してミックス精度学習の恩恵を受けるために、彼らのネットワークに少しの行の変更すらするだけです。ApexはCVPR 2018で発表されました。PyTorchコミュニティがApexのリリース以降、強力なサポートを示していることに注意してください。

Ampによって、少しの変更を行い、运行中のモデルに対して、スクリプトの作成や実行中にミックス型の問題を心配する必要がなくなります。Ampの仮定はPyTorchを普通でない方法で使用するモデルにとってもよく合わないかもしれませんが、必要であればそれらの仮定を調整するためのフックがあります。

Ampはミックス精度学習のすべての利点を提供しますが、損失拡大や型変換の明示的な管理は必要ありません。ApexのGitHubウェブサイトにはインストール手順の指示が含まれていて、公式なAPI文書はここに見つかります。

Ampの仕組み

Amphpは論理レベルでホワイトリスト/ブラックリストのパラダイムを利用しています。PyTorchのテンサー操作には、torch.nn.functional.conv2dなどのニューラルネットワーク関数、torch.logなどの単純な数学関数、torch.Tensor.add__などのテンサーメソッドが含まれます。この宇宙には3つの主要なカテゴリの関数があります:

ホワイトリスト: FP16数学のスピード向上を受けることが可能な関数。典型的な应用には、行列積算と畳み込みが含まれます。
ブラックリスト: 16ビットの精度が十分でない可能性がある関数には、FP32の入力が必要です。
その他のすべて(残った関数): FP16で実行できる関数ですが、FP32 -> FP16変換を通じてFP16で実行するためのコストが大きすぎません。

Amphpの仕事は、理論上は単純です。Amphpは、PyTorch関数がホワイトリストにあるか、ブラックリストにあるか、それ以外であるかを判定して呼び出す前に、すべての引数をFP16に変換する必要があります。引数は全て同じ型にする必要があります。このポリシーは、実際には、思い通りに適用されると思われません。

PyTorchモデルとAmpを併用する

現在のPyTorchスクリプトにAmpを含めるには、以下の手順に従います。

Apexライブラリを使用してAmpを導入します。
Ampを初期化し、モデル、最適化器、およびPyTorchの内部関数に必要な変更を行うことができます。
バックプロパゲーションの場所（.backward()）を注意しておくことで、Ampは損失をスケールし、各イテレーションの状態をクリアすることができます。

ステップ1

最初のステップには1行のコードしか必要です。

from apex import amp

ステップ2

学習に使用する神经网路モデルと最適化器はすでに指定されなければならず、このステップは1行の長さで完了します。

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

追加の設定でAmpのテンsorと操作タイプの調整を細かく調整することができます。amp.initialize()関数は多くのパラメータを受け取りますが、以下の3つのパラメータしか指定しません：

models (torch.nn.Module またはtorch.nn.Moduleのリスト) – 変更/キャストするモデル。
最適化器 (任意, torch.optim.Optimizer または torch.optim.Optimizers 的なリスト) – 修飾またはキャストする最適化器。学習には必須であり、推論には任意です。
opt_level (str, 任意であり、デフォルトは“O1”です) – 純粋なまたは混合精度和を行うための最適化レベル。受け入れられる値は“O0”、“O1”、“O2”および“O3”であり、上記で詳細に説明されています。これらは4つの最適化レベルです:

O0 は FP32 学習用: これは no-op であり、入力モデルがすでに FP32 であるため、O0 は精度の基準を設定するために役立つかもしれません。

O1 は混合精度（典型的な使用に推奨される）: Tensor と Torch のすべてのメソッドを whitelist-blacklist 入力キャストスキームを使用するよう修正する。FP16 では、Tensor Core をサポートする GEMMs やコンボルューションなどの whitelist 操作が行われます。Softmax などは FP32 精度が必要な blacklist 操作です。明記されない限り、O1 は動的な損失スケーリングをもとにしています。

O2 は “几乎FP16” の混合精度: O2 はモデルの重みを FP16 にキャストし、モデルの前向きメソッドを修飾して入力データを FP16 にキャストします。Batchnorm を FP32 に保持し、FP32 の親重みを保持し、最適化器の param_groups を更新して、最適化器.step() が直接 FP32 重みに対応して动作するようにします。O1 とは異なり、O2 は Torch の関数や Tensor のメソッドを修飾しません。

O3をFP16トレーニングに使用する：O3はO1とO2よりも真の混合精度に関する安定性を欠く可能性があり、そのため、O1とO2の効率を評価するのに基準スピードを設定するのが有利かもしれません。
O3において、keep_batchnorm_fp32=Trueという追加の属性オーバーライドを使用することは、モデルにバッチ正規化を使用している場合、光の速度のようなスピードを決定することができ、Cudnnバッチ正規化を可能にします。

O0とO3は真の混合精度ではなく、それぞれ精度とスピードの基準を設定します。混合精度の実装はO1とO2として定義されています。
O1やO2によって、特定のモデルに最も効果的なパフォーマンスと精度を向上させるのか試してみてください。

ステップ3

コードで後向傳播が行われる場所を正しく特定してください。
以下のような数行のコードが現れます：

loss = criterion(…)
loss.backward()
optimizer.step()

ステップ4

Ampコンテキストマネージャを使用することで、後向传播を包み込んだだけで損失スケーリングを有効にすることができます。

loss = criterion(…)
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
optimizer.step()

これで完了です。現在、混合精度のトレーニングを有効にしたスクリプトを再実行することができます。

関数呼び出しの捕获

PyTorchは、その柔軟性と動的性によって、静的なモデルオブジェクトやグラフを保持して、上記のようなキャストを挿入することができません。必要な機能に対して「モンキーパッチング」を行うことで、Ampは動的にパラメーターのキャストを截取ることができます。

たとえば、下記のコードを使用することで、torch.nn.functional.linearメソッドの引数が常にfp16にキャストされるようにすることができます。

orig_linear = torch.nn.functional.linear
def wrapped_linear(*args):
 casted_args = []
  for arg in args:
    if torch.is_tensor(arg) and torch.is_floating_point(arg):
      casted_args.append(torch.cast(arg, torch.float16))
    else:
      casted_args.append(arg)
  return orig_linear(*casted_args)
torch.nn.functional.linear = wrapped_linear

Ampはコードのより詳細な修飾を加えるかもしれませんが、Amp.init()を呼び出すことで、動的に正しいキャストを行うために、すべての関連するPyTorch関数にモンキーパッチを挿入することになります。

キャストの最小化

Ampはすべてのパラメーターのキャストを内部キャッシュに保持し、必要に応じて再利用します。それぞれのイテレーションで、後向きパスのコンテキストマネージャは、Ampがキャッシュをクリアする時間を示します。

PyTorchを使用した自動キャストと勾配缩放

“自動混合精度トレーニング”は、torch.cuda.amp.autocastとtorch.cuda.amp.GradScalerの組み合わせを指します。torch.cuda.amp.autocastを使用することで、特定の領域にただ単に自動キャストを設定することができます。自動キャストは、GPU操作に最適な精度を自動的に選択し、正確性を保ちながら効率を最適化します。

torch.cuda.amp.GradScalerのインスタンスは、梯度缩放の手順を簡単に実行することができます。梯度缩放は、梯度下溢を軽減し、float16の梯度を持つネットワークがより良い収束を実現するために役立ちます。

以下のコードは、autocast()を使用してPyTorchで自動混合精度を得る方法を示しています。

# デフォルトの精度でモデルと最適化器を作成
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# トレーニングの初めに一度GradScalerを作成します。
scaler = GradScaler()

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

        # autocastで前向きパスを実行します。
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)

        # 後向き演算は、autocastが選んだ同じ型で実行されます。
        scaler.scale(loss).backward()

        # scaler.step()は最初に最適化器の割り当てられたパラメータの勾配を未缩放状態に戻します。
   
        scaler.step(optimizer)

        # 次の iteration に更新するスケールを更新します。
        scaler.update()

特定の操作の前向きパスにfloat16の入力がある場合、その操作の後向きパスはfloat16の勾配を生成し、float16は小さな大小の勾配を表現できないかもしれません。

これらの値が0に書き換えられる(“下位溢出”)ことがあるため、関連パラメータの更新が失われる可能性があります。

勾配スケーリングは、スケール因子を使用してネットワークの損失を倍增し、スケールされた損失に基づいて後向きパスを行い、下位溢出を防止する技術です。同じ因子でネットワーク内で後向き勾配をスケールする必要があります。結果として、勾配の大きさが大きくなり、0に書き換えられることが防止されます。

パラメータ更新前に、各パラメータの勾配(.grad属性)にスケール因子を適用する必要があります。これは学習率に影響を与えないようにするためです。autocastとGradScalerはモジュール的なため、独立して使用することができます。

未スケール勾配での操作

勾配クリッピング

すべての勾配は、Scaler.scale(Loss).backward() メソッドを使用してスケーリングできます。backward() と scaler.step(optimizer) の間のパラメータの .grad プロパティは、変更または確認する前にスケーリング解除されなければなりません。グラデーションセットのグローバルノルム（torch.nn.utils.clip_grad_norm_()参照）または最大の大きさ（torch.nn.utils.clip_grad_value_()参照）を特定の値（ユーザーによって設定されたしきい値）以下に制限したい場合、”勾配クリッピング” と呼ばれる技術を使用できます。

スケーリング解除せずにクリッピングすると、勾配のノルムや最大値がスケーリングされ、要求されたしきい値（スケーリング解除された勾配のためのしきい値であるはずだったもの）が無効になります。オプティマイザーのパラメータに含まれる勾配は、scaler.unscale(optimizer) によってスケーリング解除されます。
他のオプティマイザー（たとえば optimizer1）に以前割り当てられたパラメータの勾配をスケーリング解除するには、scaler.unscale(optimizer1) を使用できます。この概念を示すために、2行のコードを追加して説明できます。

# オプティマイザーに割り当てられたパラメータの勾配をその場でスケーリング解除
        scaler.unscale_(optimizer)
# オプティマイザーに割り当てられたパラメータの勾配がスケーリング解除されたので、通常通りクリッピング: 
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

スケールされたグラデーションと共働く

グラデーション積み重ね

グラデーション積み重ねは非常に基本な概念に基づいています。モデルのパラメーターを更新する代わりに、損失とグラデーションを計算するために连续的なバッチを通してグラデーションを積み重ねることができます。

特定の数のバッチ後、集計されたグラデーションに基づいてパラメーターが更新されます。以下は、グラデーション積み重ねを pytorchで使用するコードの一部です。

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
            # 損失を正規化 
            loss = loss / iters_to_accumulate

        # スケールされたグラデーションを積み重ねる。
        scaler.scale(loss).backward()
          # 重みの更新
        if (i + 1) % iters_to_accumulate == 0:
            # 必要があればここでunscale_を実行します 
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

グラデーション積み重ねは、batch_per_iter * iters_to_accumulateの適切なバッチサイズでグラデーションを積み重ねます。
スケールを有効なバッチに調整することが重要です。これには無限大/NaN値のグラデーションの確認、無限大/NaN値が発見された場合のステップをスキップし、有効なバッチの粒度によってスケールを更新する必要があります。
特定の有効なバッチのグラデーションを加算する際、グラデーションをスケールされた一致したスケール因子で保持することも非常に重要です。

もし修士の学習率が未調整（または調整係数が変更され）て完了する前に、次の逆向きパスは調整された学習率を未調整の学習率（または異なる係数で調整された学習率）に加えるため、积累された未調整の学習率のステップを復元することはできなくなります。

修士はunscale学習率を、次のステップのための調整された学習率が积累された後、すぐにステップの前にunscaleを使用して可能です。
効果的なバッチを保証するために、各迭代の終了時にupdateを呼び出します。ここで以前にstepを呼び出した場合です。
enumerate(data)関数は、データをイテレートしながらバッチのインデックスを追跡することができます。
進行中の損失をiters_to_accumulate(loss / iters_to_accumulate)で割ることで、損失を正規化して処理している各ミニバッチの貢献を減らすことができます。各バッチ内で損失を平均化すると、割り当てはすでに正しく、さらなる正規化は不要です。このステップは、損失を計算する方法によって必要でないかもしれません。
scaler.scale(loss).backward()を使用すると、PyTorchは缩放された勾配を积累し、optimizer.zero_grad()を呼び出さない限りそれらを保持します。

勾配ペナルティ

勾配ペナルティを実装する際には、torch.autograd.grad()を使用して勾配を構築し、これらを結合してペナルティ値を形成し、その後損失に加えます。缩放生ずれなしや自動キャストなしのL2ペナルティの例が以下に示されています。

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)

        # 勾配を作成
        grad_prams = torch.autograd.grad(outputs=loss,
                                          inputs=model.parameters(),
                                          create_graph=True)

        # ペナルティ項を計算し、損失に加えます
        grad_norm = 0
        for grad in grad_prams:
            grad_norm += grad.pow(2).sum()
        grad_norm = grad_norm.sqrt()
        loss = loss + grad_norm

        loss.backward()

        # ここで勾配をクリップすることができます

        optimizer.step()

torch.autograd.grad()に与えられるテンサーは、勾配ペナルティを実装するために缩放されなければならない。勾配を結合してペナルティ値を得る前に、勾配を解除缩放する必要があります。勾配項の計算は前向きパスの一部であるため、自動キャストコンテキストの内部で行われなければならない。
同じL2ペナルティについて、以下のように表示されます。

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # autograd.gradの後向きパスにて損失拡大を行います。#scaled_grad_prams
        scaled_grad_prams = torch.autograd.grad(outputs=scaler.scale(loss),
                                                 inputs=model.parameters(),
                                                 create_graph=True)

        # 罰則を計算する前にgrad_pramsを作成します（grad_pramsは#缩放されていない必要があります）。 
        # オプティマイザーがscaled_grad_pramsを所有していないため、scaler.unscale_を使用せずに通常の除法を行います。
        inv_scaled = 1./scaler.get_scale()
        grad_prams = [p * inv_scaled for p in scaled_grad_prams]

        # 罰則項を計算して損失に加えます。 
        with autocast():
            grad_norm = 0
            for grad in grad_prams:
                grad_norm += grad.pow(2).sum()
            grad_norm = grad_norm.sqrt()
            loss = loss + grad_norm

        # 後向き呼び出しにて缩放を適用します。
        # 適切に缩放された叶の勾配を積みます。
        scaler.scale(loss).backward()

        # ここでunscale_することができます 

        # step()とupdate()は通常通り続けます。
        scaler.step(optimizer)
        scaler.update()

複数のモデル、損失、およびオプティマイザーでの作業

Scaler.scaleは、ネットワーク内に多くの損失がある場合、各損失に対して呼び出されなければならない。
ネットワーク内に多くのオプティマイザーがある場合、scaler.unscaleを其中の任何一个に実行することができ、各オプティマイザーに対してscaler.stepを呼び出す必要がある。しかし、scaler.updateは、このイテレーションで使用されるすべてのオプティマイザーのステップ後に一度だけ使用するべきです：

scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer1.zero_grad()
        optimizer2.zero_grad()
        with autocast():
            output1 = model1(input)
            output2 = model2(input)
            loss1 = loss_fn(2 * output1 + 3 * output2, target)
            loss2 = loss_fn(3 * output1 - 5 * output2, target)

       # retain graphはAMPに関連していないが、この例では後退()コールが共有するグラフの特定の領域に関連しています。 
        scaler.scale(loss1).backward(retain_graph=True)
        scaler.scale(loss2).backward()

        # paramsの梯度を表示または調整したい場合、明示的にアンスケールされるオプティマイザーを指定することができます。 .
        scaler.unscale_(optimizer1)

        scaler.step(optimizer1)
        scaler.step(optimizer2)

        scaler.update()

各オプティマイザーは、自身の梯度に対してinfs/NaNsを確認し、ステップをスキップするか否か個别的に判断します。一部のオプティマイザーはステップをスキップすることがあり、他のオプティマイザーはそうでないかもしれません。ステップスキップは数百回のイテレーションごとに一度だけ発生するので、 convergenceに影響を与えてはならない。マルチオプティマイザーモデルについて、梯度スケーリングを追加した後に劣化した convergenceを見つけた場合、問題を報告することができます。

複数GPUでの作業

深層学習モデルの最も重要な問題の1つは、1つのGPUでトレーニングするには过大きくなり、1つのGPUでモデルをトレーニングする時間が長すぎるためです。有名な研究者は、ImageNetのトレーニング周期を2週間から18分に短縮し、最も大きく、最も進歩したTransformer-XLを4年間でなく2週間でトレーニングすることができました。

DataParallelとDistributedDataParallel

品質をcompromiseすることなく、PyTorchは使用の簡便さとコントロールの最好的な組み合わせを提供します。nn.DataParallelとnn.parallel.DistributedDataParallelは、PyTorchの多くのGPUでトレーニングを行うための2つの機能です。これらの簡単に使用できるwrapperを使用し、GPU間での変更を行い、複数のGPUでネットワークを学習させることができます。

单 processes 中的 DataParallel

一台のマシンで、DataParallelはGPU間でトレーニングを分散させることができます。
DataParallelが本当にどのように機能するのかを詳しく見てみましょう。
DataParallelを使用してニューラルネットワークの学習を行う際には、以下の段階が行われます。

Source:

Mini-batchがGPU:0に分割されます。
全ての利用可能なGPUに分割され、配布されます。
モデルをGPUsにコピーします。
すべてのGPUで前向きパスが行われます。
GPU:0でネットワークの出力に関連する損失を計算し、それらを各GPUに返信します。各GPUで勾配を計算する必要があります。
GPU:0上で勾配の和を計算し、最適化器を適用してモデルを更新する。

注意すべきことは、ここで讨论されている懸念はautocastによってのみ適用されることに限られています。GradScalerの動作は変更していません。torch.nn.DataParallelが各デバイスに対してスレッドを作成して前向きパスを実行するかどうかは、重要ではありません。autocast状態は各々でやりとりされ、以下が機能するでしょう。

model = Model_m()
p_model = nn.DataParallel(model)

# main threadでautocastを設定
with autocast():
    # p_modelにもautocastingが行われます 
    output = p_model(input)
    # loss_fnもautocast
    loss = loss_fn(output)

プロセスごとに1GPUのDistributedDataParallel

torch.nn.parallel.DistributedDataParallelのドキュメントでは、最適な性能を得るために、每个プロセスに1つのGPUを使用することを推奨しています。この状況で、DistributedDataParallelは内部でスレッドを起動しませんので、autocastとGradScalerの使用は影響を受けません。

DistributedDataParallel, プロセスに複数のGPUを使用する

ここで、torch.nn.parallel.DistributedDataParallelは、torch.nn.DataParallelのように各デバイス上で前向きパスを実行するためにサイドスレッドを生成することがあります。解決法は同じです: あなたのモデルの前向きメソッドの一部としてautocastを適用することで、サイドスレッドで有効にすることを Ensure します。

結論

この記事では、以下を説明しました:

Apexを導入しました。
Ampsの仕組みを見ました。
勾配スケーリング、勾配クリッピング、勾配積积累と勾配ペナルティを実施する方法を見ました。
複数のモデル、損失関数、最適化器と共に工作する方法を見ました。
複数のGPUで1つのプロセスでDataParallelを実行する方法を見ました。

参考文献

https://developer.nvidia.com/blog/apex-pytorch-easy-mixed-precision-training/
https://nvidia.github.io/apex/amp.html
https://discuss.pytorch.org/t/accumulating-gradients/30020
https://towardsdatascience.com/how-to-scale-training-on-multiple-gpus-dae1041f49d2

Source:
https://www.digitalocean.com/community/tutorials/automatic-mixed-precision-using-pytorch