[モデルの高速展開] - Pytorch 自動混合精度トレーニング

自動混合精度

torch.amp は、混合精度の便利なメソッドを提供します。一部の操作では torch.float32 (float) データ型が使用され、その他の操作では低精度の float データ型 ( lower_precision_fp )、つまり torch.float16 (half) または torch .bfloat16 が使用されます。線形層や畳み込みなどの一部の操作は、 lower_precision_fp の方がはるかに高速です。ダウンスケーリングなどの他の操作では、多くの場合、float32 のダイナミック レンジが必要になります。混合精度では、各操作を適切なデータ型に一致させようとします。

通常、データ型 torch.float16 の「自動混合精度トレーニング」は、torch.autocast と torch.cuda.amp.GradScaler を一緒に使用します。torch.autocast と torch.cuda.amp.GradScaler はモジュール式であり、必要に応じて個別に使用できます。

CUDA と CPU の場合、API も個別に提供されます。

torch.autocast("cuda", args...) は torch.cuda.amp.autocast(args...) と同等です。

torch.autocast("cpu", args...) は torch.cpu.amp.autocast(args...) と同等です。CPU の場合、現在、torch.bfloat16 の低精度浮動小数点データ型のみがサポートされています。

自動変換（オートキャスト）

torch.autocast( device_type、dtype=None、enabled=True、cache_enabled=None )

パラメータ:
device_type (str、必須) - 'cuda' または 'cpu' を使用 有効なデバイス
(bool、オプション) - ゾーンでオートキャストを有効にするかどうか。デフォルト: True
dtype (torch_dtype、オプション) – torch.float16 または torch.bfloat16 を使用するかどうか。
cache_enabled (ブール値、オプション) - 自動キャスト内部重みキャッシュを有効にするかどうか。デフォルト: True

autocast はコンテキスト マネージャーまたはデコレーターとして使用でき、スクリプト スコープを混合精度で実行できます。

これらのリージョンでは、データ操作は、精度を維持しながらパフォーマンスを向上させるために、その操作に固有の自動キャストによって選択された dtype で実行されます。

オートキャストが有効な領域に入るとき、テンソルは任意のタイプにすることができます。オートキャストを使用する場合は、モデルまたは入力に対してhalf() または bfloat16() を呼び出さないでください。オートキャストは、損失計算を含むネットワークの順方向パスのみをラップする必要があります。自動キャストでのバックパスは推奨されません。逆方向操作では、自動キャストが対応する順方向操作に使用するのと同じタイプの実行が行われます。

CUDA の例:

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

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

    # Enables autocasting for the forward pass (model + loss)
    with autocast():
        output = model(input)
        loss = loss_fn(output, target)

    # Exits the context manager before backward()
    loss.backward()
    optimizer.step()

結果として得られる浮動小数点テンソルは、自動変換が有効になっている領域では float16 になる場合があります。異なる dtype の浮動小数点テンソルでこれらを使用すると、自動変換が無効になっている領域に戻った後に型不一致エラーが発生する可能性があります。その場合は、自動変換ゾーンで得られたテンソルを float32 (または必要に応じて他の dtype) に変換して戻します。自動変換された領域のテンソルがすでに float32 である場合、変換は何も行われず、追加のオーバーヘッドは発生しません。

CUDA の例:

# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")

with autocast():
    # torch.mm is on autocast's list of ops that should run in float16.
    # Inputs are float32, but the op runs in float16 and produces float16 output.
    # No manual casts are required.
    e_float16 = torch.mm(a_float32, b_float32)
    # Also handles mixed input types
    f_float16 = torch.mm(d_float32, e_float16)

# After exiting autocast, calls f_float16.float() to use with d_float32
g_float32 = torch.mm(d_float32, f_float16.float())

CPU トレーニングの例:

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

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

        # Runs the forward pass with autocasting.
        with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
            output = model(input)
            loss = loss_fn(output, target)

        loss.backward()
        optimizer.step()

CPU 推論の例:

# Creates model in default precision
model = Net().eval()

with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    for input in data:
        # Runs the forward pass with autocasting.
        output = model(input)

jit トレースを使用した CPU の推論例:

class TestModel(nn.Module):
    def __init__(self, input_size, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(input_size, num_classes)
    def forward(self, x):
        return self.fc1(x)

input_size = 2
num_classes = 2
model = TestModel(input_size, num_classes).eval()

# For now, we suggest to disable the Jit Autocast Pass,
# As the issue: https://github.com/pytorch/pytorch/issues/75956
torch._C._jit_set_autocast_mode(False)

with torch.cpu.amp.autocast(cache_enabled=False):
    model = torch.jit.trace(model, torch.randn(1, input_size))
model = torch.jit.freeze(model)
# Models Run
for _ in range(3):
    model(torch.randn(1, input_size))

autocast(enabled=False) サブリージョンは、オートキャストが有効なリージョン内にネストできます。自動キャストを部分的に無効にすると、特定のデータ型でサブ領域を強制する場合に便利です。オートキャストを無効にすると、実行タイプを明示的に制御できるようになります。サブ領域では、周囲の領域からの入力を使用する前に指定されたデータ型に変換する必要があります。

# 在默认数据类型（这里假设为float32）中创建一些张量
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")

with autocast():
    e_float16 = torch.mm(a_float32, b_float32)
    with autocast(enabled=False):
        # 调用e_float16.float()以确保使用float32执行
        # （这是必需的，因为e_float16是在autocast区域中创建的）
        f_float32 = torch.mm(c_float32, e_float16.float())

    # 当重新进入启用autocast的区域时，无需手动转换类型。
    # torch.mm仍然以float16运行并产生float16的输出，不受输入类型的影响。
    g_float16 = torch.mm(d_float32, f_float32)

オートキャスト状態はスレッドローカルです。新しいスレッドでこれを有効にしたい場合は、そのスレッドでコンテキスト マネージャーまたはデコレーターを呼び出す必要があります。

torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)
torch.cuda.amp.autocast(args…)等と同じtorch.autocast(“cuda”, args…)

* torch.cuda.amp.custom_fwd(fwd=None, , Cast_inputs=None)
カスタム autograd 関数の forward メソッドの補助デコレータ (torch.autograd.Function のサブクラス)。

パラメータ:
Cast_inputs (torch.dtype または None、オプション、default=None) - None でない場合、オートキャストが有効な領域で順方向に実行するときに、受信浮動小数点 CUDA テンソルをターゲット データ型に変換します (非浮動小数点テンソルは影響を受けません)、オートキャストを無効にして転送を実行します。None の場合、現在の自動キャスト状態に従って転送の内部操作が実行されます。

勾配スケーリング

演算の順方向パスに float16 入力がある場合、その演算の逆方向パスは float16 の勾配を生成します。大きさが小さい勾配値は float16 として表現できない場合があります。これらの値はゼロにされる (「アンダーフロー」) ため、対応するパラメーターへの更新は失われます。

アンダーフローを防ぐために、「勾配スケーリング」は、ネットワークの損失にスケーリング係数を乗算し、スケーリングされた損失に対してバック パスを実行することによって機能します。ネットワークを介して伝播される勾配も同じ係数でスケールされます。言い換えれば、勾配値の大きさは大きいため、ゼロ化されません。

スケーリング係数が学習率の設定に干渉しないように、オプティマイザーがパラメーターを更新する前に、各パラメーターの勾配 (.grad プロパティ) を復元する必要があります。

torch.cuda.amp.GradScaler(init_scale=65536.0、growth_factor=2.0、backoff_factor=0.5、growth_interval=2000、enabled=True)

get_backoff_factor() は、スケーリング バックオフ係数を含む Python float を返します。

get_growth_factor() は、スケーリング成長係数を含む Python float を返します。

get_growth_interval() は、成長間隔を含む Python 整数を返します。

get_scale() は、現在のスケーリング係数を含む Python float を返します。スケーリングが無効な場合は 1.0 を返します。

**警告:** get_scale() は CPU-GPU 同期を生成します。

is_enabled() は、このインスタンスが有効かどうかを示すブール値を返します。

load_state_dict(state_dict) ロード スケーラーの状態。このインスタンスが無効になっている場合、load_state_dict() は何も行いません。

パラメータ: state_dict (dict) – スケーラーの状態。state_dict() を呼び出して返されるオブジェクトである必要があります。

scale(outputs) は、スケール係数によってテンソルまたはテンソルのリストをスケーリングします。

スケーリングされた出力を返します。GradScaler のインスタンスが有効になっていない場合は、未変更の出力を返します。

パラメータ: 出力 (Tensor または Tensor の反復可能) – スケーリングされた出力。

set_backoff_factor(new_factor) パラメータ: new_factor (float) – 新しいスケーリング バックオフ係数として使用する値。

set_growth_factor(new_factor) パラメータ: new_factor (float) – 新しいスケーリング成長係数として使用する値。

set_growth_interval(new_interval) パラメータ: new_interval (int) – 新しい成長間隔として使用する値。

state_dict() は、スケーラーの状態を辞書として返します。これには 5 つのエントリが含まれています。

"scale" - 現在のスケールを含む Python float

"growth_factor" - 現在の成長係数を含む Python float

"backoff_factor" - 現在のバックオフ係数を含む Python float

"growth_interval" - 現在の成長間隔を含む Python 整数

「_growth_tracker」 - 最新の連続したスキップされていないステップの数を含む Python 整数。

このインスタンスが有効になっていない場合は、空の辞書を返します。

注:特定の反復後にポイント スケーラーの状態を確認したい場合は、update() の後に state_dict() を呼び出す必要があります。

step(オプティマイザー、*args、**kwargs)

step() は次の 2 つの操作を実行します。

1 は、内部的に unscale_(optimizer) を呼び出します (反復の前に unscale_() が明示的に呼び出されていない限り)。unscale_() の一部として、勾配に inf/NaN が含まれているかどうかがチェックされます。

2 inf/NaN 勾配が見つからない場合、optimizer.step() はスケールされていない勾配を使用して呼び出されます。それ以外の場合、パラメータの破損を避けるために optimizer.step() がスキップされます。

*args と **kwargs は optimizer.step() に渡されます。

optimizer.step(*args, **kwargs) の戻り値を返します。

パラメータ: optimizer (torch.optim.Optimizer) – グラデーションを適用するオプティマイザー。

args – 任意の引数。

kwargs – 任意のキーワード引数。

警告する

クロージャの使用は現在サポートされていません。

unscale_(optimizer) は、オプティマイザーの勾配テンソルをスケーリング係数で除算 (「スケール解除」) します。

unscale_() はオプションであり、バックプロパゲーションとステップ ( step() ) の間で勾配が変更または検査される場合に適しています。unscale_() が明示的に呼び出されない場合、グラデーションはステップ ( step() ) 中に自動的にスケール解除されます。

unscale_() を使用してスケールされていないグラデーションのクリッピングを有効にする簡単な例:

…scaler.scale(loss).backward()scaler.unscale_(optimizer)torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)scaler.step(optimizer)scaler.update()パラメータ: optimizer (torch. optim.Optimizer) – スケールを解除するグラデーションを備えたオプティマイザー。

知らせ

unscale_() は CPU-GPU 同期を生成しません。

警告する

unscale_() への各呼び出しは、オプティマイザーごとの step() 呼び出しごとに 1 回だけ呼び出す必要があり、そのオプティマイザーに割り当てられたパラメーターに対してすべての勾配が累積された後にのみ呼び出す必要があります。各 step() の間に特定のオプティマイザーに対して unscale_() を 2 回連続して呼び出すと、RuntimeError が発生します。

警告する

unscale_() は、.grad 属性を置き換えて、不可逆的な方法でスパース グラデーションのスケールを解除する場合があります。

update(new_scale=None)[SOURCE] スケーリング係数を更新します。

オプティマイザーのステップがスキップされた場合、そのステップは、backoff_factor とスケーリング係数を乗算した分だけ削減されます。スキップされていない Groove_interval の反復が連続してある場合は、growth_factor にスケーリング係数を乗算して増加させます。

new_scale を渡すと、新しいスケール値が手動で設定されます。(new_scale は直接使用されませんが、GradScaler の内部スケーリング テンソルを設定するために使用されます。したがって、new_scale がテンソルの場合、そのテンソルに対するインプレース変更は、GradScaler によって内部的に使用されるスケールにそれ以上の影響を与えません。)

パラメータ: new_scale (float または torch.cuda.FloatTensor、オプション、デフォルトは None) – 新しいスケーリング係数。

警告する

update() は、使用されるすべてのオプティマイザーに対してscaler.step(optimizer) が呼び出された反復の最後にのみ呼び出す必要があります。

Autocast Op 関連のリファレンス

オートキャスト オペレーションの適格性

オートキャストが有効かどうかに関係なく、float64 または非 float 型で動作する操作はオートキャストされず、元の型で実行されます。autocast は、アウトオブプレイス操作と tensor メソッドにのみ影響します。out=...tensor を明示的に提供するインプレース操作と呼び出しは、オートキャストが有効なリージョンで許可されますが、オートキャストは通過しません。たとえば、自動遷移が有効になっているリージョンでは、a.addmm(b,c) は自動変換できますが、a.addmm_(b,c,out=d) は自動変換できません。最高のパフォーマンスと安定性を得るには、自動キャストが有効なリージョンでのアウトオブプレース操作を優先してください。dtype=... を明示的に呼び出す操作はオートキャストの使用には適格ではなく、dtype 引数の出力を生成します。

CUDA Op 固有の動作

次のリストでは、自動移行が有効になっているリージョンでの対象となる操作の動作について説明します。これらの操作は、torch.nn.Module の一部として呼び出されるか、関数として呼び出されるか、torch.Tensor メソッドとして呼び出されるかに関係なく、常に自動的に変換されます。関数が複数の名前空間で公開されている場合、それらは名前空間に関係なく自動的に変換されます。

以下にリストされていない操作は自動変換されません。これらは、入力によって定義されたタイプに従って動作します。ただし、リストに記載されていない操作が自動変換された操作の下流にある場合、自動変換によっても、その操作によって実行される操作のタイプが変更される可能性があります。

演算がリストされていない場合、その演算は float16 で数値的に安定していると想定されます。リストされていない演算が float16 で数値的に不安定であると思われる場合は、問題を報告してください。

float16に自動変換できるCUDA Ops

matmul、addbmm、addmm、addmv、addr、baddbmm、bmm、chain_matmul、multi_dot、conv1d、conv2d、conv3d、conv_transpose1d、conv_transpose2d、conv_transpose3d、GRUCell、linear、LSTMCell、matmul、mm、mv、prelu、RNNCell

float32に自動変換できるCUDA Ops

pow、rdiv、rpow、rtruediv、acos、asin、binary_cross_entropy_with_logits、cosh、cosine_embedding_loss、cdist、cosine_similarity、cross_entropy、cumprod、cumsum、dist、erfinv、exp、expm1、group_norm、hinge_embedding_loss、kl_div、l1_loss、layer_norm、log、log_softmax 、 log10、log1p、log2、margin_ranking_loss、mse_loss、multilabel_margin_loss、multi_margin_loss、nll_loss、norm、normalize、pdist、poisson_nll_loss、pow、prod、reciprocal、rsqrt、sinh、smooth_l1_loss、soft_margin_loss、softmax、softmin、softplus、sum、renorm、tan、トリプレットマージンロス

最も幅広い入力タイプにブーストできる CUDA Ops

これらの操作では、安定性のために特定のデータ型は必要ありませんが、複数の入力が必要であり、入力のデータ型が一致する必要があります。すべての入力が float16 の場合、演算は float16 で動作します。いずれかの入力が float32 である場合、autoconvert はすべての入力を float32 に変換し、float32 で演算を実行します。

addcdiv、addcmul、atan2、バイリニア、クロス、ドット、grid_sample、index_put、scatter_add、tensordot

ここにリストされていない一部の操作 (たとえば、add などのバイナリ操作) は、自動変換を介さずに、それ自体で入力を促進できます。入力が bfloat16 と float32 の混合である場合、自動変換が有効かどうかに関係なく、これらの演算は float32 で動作し、float32 出力が生成されます。

binary_cross_entropy よりも binary_cross_entropy_with_logits を優先します

torch.nn.Functional.binary_cross_entropy() (およびそれをラップする torch.nn.BCELoss) の逆伝播により、float16 では表現できない勾配が生成される可能性があります。自動変換が有効な領域では、順方向入力は float16 である可能性があります。これは、バックプロパゲーションの勾配が float16 で表現できる必要があることを意味します (バックプロパゲーションでは変換が逆になる必要があるため、float16 順方向入力を float32 に自動変換することは役に立ちません)。したがって、binary_cross_entropy と BCELoss は、自動遷移が有効になっている領域でエラーをスローします。

多くのモデルは、バイナリ クロスエントロピー層の前にシグモイド層を使用します。この場合、torch.nn.functioner.binary_cross_entropy_with_logits() または torch.nn.BCEWithLogitsLoss を使用して 2 つのレイヤーを結合します。binary_cross_entropy_with_logits および BCEWithLogits は安全に自動変換できます。

CPU 操作固有の動作

次のリストでは、自動移行が有効になっているリージョンでの対象となる操作の動作について説明します。これらの操作は、torch.nn.Module の一部として呼び出されるか、関数として呼び出されるか、torch.Tensor メソッドとして呼び出されるかに関係なく、常に自動的に変換されます。関数が複数の名前空間で公開されている場合、それらは名前空間に関係なく自動的に変換されます。

以下にリストされていない操作は自動変換されません。これらは、入力によって定義されたタイプに従って動作します。ただし、リストに記載されていない操作が自動変換された操作の下流にある場合、自動変換によっても、その操作によって実行される操作のタイプが変更される可能性があります。

演算がリストされていない場合、その演算は bfloat16 で数値的に安定していると想定されます。リストされていない演算が bfloat16 で数値的に不安定であると思われる場合は、問題を報告してください。

bfloat16 に自動的に変換できる CPU Ops

conv1d、conv2d、conv3d、bmm、mm、baddbmm、addmm、addbmm、linear、matmul、_convolution

float32 に自動変換できる CPU

作戦conv_transpose1d、conv_transpose2d、conv_transpose3d、avg_pool3d、binary_cross_entropy、grid_sampler、grid_sampler_2d、_grid_sampler_2d_cpu_fallback、grid_sampler_3d、polar、prod、quantile、nanquantile、stft、cdist、trace、view_as_complex、cholesky、chole sky_inverse、cholesky_solve、inverse、lu_solve、orgqr、inverse、ormqr、 pinverse、max_pool3d、max_unpool2d、max_unpool3d、adaptive_avg_pool3d、reflection_pad1d、reflection_pad2d、replication_pad1d、replication_pad2d、replication_pad3d、mse_loss、ctc_loss、kl_div、multilabel_margin_loss、fft_fft、fft_ifft、fft_fft2、fft_ifft2、 fft_fftn、fft_ifftn、fft_rfft、fft_irfft、fft_rfft2、fft_irfft2、fft_rfftn、 fft_irfftn、fft_hfft、fft_ihfft、linalg_matrix_norm、linalg_cond、linalg_matrix_rank、linalg_solve、linalg_cholesky、linalg_svdvals、linalg_eigvals、linalg_eigvalsh、linalg_inv、linalg_householder_product、linalg_tensorinv、linalg_ tensorsolve、fake_quantize_per_tensor_affine、eig、geqrf、lstsq、_lu_with_info、qr、solve、svd、symeig、triangular_solve、fractional_max_pool2d、fractional_max_pool3d、adaptive_max_pool3d、multilabel_margin_loss_forward、linalg_qr、linalg_cholesky_ex、linalg_svd、linalg_eig、linalg_eigh、 linalg_lstsq、linalg_inv_ex

最も幅広い入力タイプにブーストできる CPU Ops

これらの操作では、安定性のために特定のデータ型は必要ありませんが、複数の入力が必要であり、入力のデータ型が一致する必要があります。すべての入力が bfloat16 の場合、演算は bfloat16 で動作します。いずれかの入力が float32 の場合、autoconvert はすべての入力を float32 に変換し、float32 で演算を実行します。

猫、スタック、インデックスコピー

ここにリストされていない一部の操作 (たとえば、add などのバイナリ操作) は、自動変換を介さずに、それ自体で入力を促進できます。入力が bfloat16 と float32 の混合である場合、自動変換が有効かどうかに関係なく、これらの演算は float32 で動作し、float32 出力が生成されます。

補足：inplaceとout of inplaceの理解（inplaceを理解して理解する）

inplace=True は、インプレース操作を実行することを意味し、インプレース カバレッジ操作を実行することを選択します。たとえば、x+=1 は、元の値 x を演算し、得られた結果でその値を直接上書きします。y=x+5、x=y は x に対するインプレース演算ではありません。
inplace=True 操作の利点は、コンピューティング メモリを節約し、他の無関係な変数を節約できることです。
注: inplace=True を使用すると、上位ネットワークから渡されたテンソルが直接変更されて入力データが変更されます。具体的な意味は次の例に示されています。

import torch
import torch.nn as nn

relu = nn.ReLU(inplace=True)
input = torch.randn(7)

print("输入数据：",input)

output = relu(input)
print("ReLU输出：", output)

print("ReLU处理后，输入数据：")
print(input)

torch.autograd.gradfunction は、PyTorch で勾配を計算するために使用される関数の 1 つです。これは、一連の変数に関する 1 つ以上のスカラー関数の勾配を計算するために使用されます。

関数のシグネチャは次のとおりです。

mathematicaCopy code
torch.autograd.grad(outputs, inputs, grad_outputs=None, retain_graph=None, create_graph=False, only_inputs=True, allow_unused=False)

パラメータの説明:

• outputs: 勾配を計算する必要があるスカラー関数を含むテンソルまたはテンソルのリスト。
• inputs: 勾配を計算する必要がある変数のテンソルまたはテンソルのリスト。
• grad_outputs:outputs勾配を計算するときに外側の勾配を指定するために使用されるテンソル、または同じ形状のテンソルのリスト。デフォルトは None で、単位勾配 (つまり 1) を使用することを意味します。
• retain_graph: ブール値。勾配の計算後に後続の計算のために計算グラフを保持するかどうかを指定します。デフォルトは「なし」です。これは、計算グラフを保持するかどうかを自動的に判断することを意味します。
• create_graph: ブール値。高次導関数を計算するための新しい計算グラフを作成するかどうかを指定します。デフォルトは False です。
• only_inputs: ブール値、入力の勾配のみを計算するかどうかを指定します。デフォルトは True で、入力の勾配のみが計算されることを意味します。
• allow_unused: 勾配を計算するときに未使用の入力を許可するかどうかを指定するブール値。デフォルトは False で、未使用の入力は許可されないことを意味します。

inputsこの関数は、 と同じ形状を持つテンソルまたはテンソルのリストを返し、inputsに対する相対的な勾配を表します。入力に勾配が必要ない場合、対応する位置の勾配は「なし」になります。

使用例を次に示します。

pythonCopy codeimport torch

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
grads = torch.autograd.grad(y, x)

print(grads)  # 输出 [tensor([4.])]

y = x ** 2上記の例では、に関する勾配を計算しx、torch.autograd.grad関数を通じて結果を取得しました。この例では、grads値 4.0 は、y相対的な勾配がx4.0 であることを意味します。