自動混合精度套件 - torch.amp

torch.amp 提供了混合精度的便利方法，其中某些運算使用 torch.float32 (float) 資料類型，而其他運算則使用較低精度的浮點資料類型 (lower_precision_fp)：torch.float16 (half) 或 torch.bfloat16。 某些運算，例如線性層和卷積，在 lower_precision_fp 中速度更快。 其他運算，例如縮減，通常需要 float32 的動態範圍。 混合精度嘗試將每個運算匹配到其適當的資料類型。

通常，使用 torch.float16 資料類型的「自動混合精度訓練」會同時使用 torch.autocast 和 torch.amp.GradScaler，如 自動混合精度範例 和 自動混合精度配方 中所示。 但是，torch.autocast 和 torch.GradScaler 是模組化的，如果需要，可以單獨使用。 如 torch.autocast 的 CPU 範例部分所示，在 CPU 上使用 torch.bfloat16 資料類型的「自動混合精度訓練/推論」僅使用 torch.autocast。

警告

torch.cuda.amp.autocast(args...) 和 torch.cpu.amp.autocast(args...) 將被棄用。 請改用 torch.autocast("cuda", args...) 或 torch.autocast("cpu", args...)。torch.cuda.amp.GradScaler(args...) 和 torch.cpu.amp.GradScaler(args...) 將被棄用。 請改用 torch.GradScaler("cuda", args...) 或 torch.GradScaler("cpu", args...)。

torch.autocast 和 torch.cpu.amp.autocast 是 1.10 版的新功能。

自動轉換

torch.amp.autocast_mode.is_autocast_available(device_type)

傳回一個布林值，指示在 device_type 上是否可以使用自動轉換。

參數

device_type (str) – 要使用的設備類型。 可能的值為：'cuda'、'cpu'、'xpu' 等。 該類型與 torch.device 的 type 屬性相同。 因此，您可以使用 Tensor.device.type 獲取張量的設備類型。

回傳類型

bool

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

autocast 的實例充當上下文管理器或裝飾器，允許您的腳本區域以混合精度運行。

在這些區域中，運算以自動轉換選擇的運算特定 dtype 運行，以提高效能，同時保持準確性。 有關詳細資訊，請參閱自動轉換運算參考。

進入啟用自動轉換的區域時，張量可以是任何類型。 使用自動轉換時，您不應在您的模型或輸入上呼叫 half() 或 bfloat16()。

autocast 應僅包裝您網路的前向傳遞，包括損失計算。 不建議在自動轉換下進行反向傳遞。 反向運算以自動轉換用於相應前向運算的相同類型運行。

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 torch.autocast(device_type="cuda"):
        output = model(input)
        loss = loss_fn(output, target)

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

請參閱自動混合精度範例，了解在更複雜的場景（例如，梯度懲罰、多個模型/損失、自定義 autograd 函數）中的用法（以及梯度縮放）。

autocast 也可以用作裝飾器，例如，在您模型的 forward 方法上

class AutocastModel(nn.Module):
    ...
    @torch.autocast(device_type="cuda")
    def forward(self, input):
        ...

在啟用 autocast 的區域中產生的浮點 Tensor 可能是 float16。返回到禁用 autocast 的區域後，將它們與不同 dtypes 的浮點 Tensor 一起使用可能會導致類型不匹配錯誤。如果是這樣，請將在 autocast 區域中產生的 Tensor 轉換回 float32（或如果需要，轉換為其他 dtype）。如果來自 autocast 區域的 Tensor 已經是 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 torch.autocast(device_type="cuda"):
    # 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 Trace 的 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 的區域中的類型不匹配錯誤是一個 bug；如果您觀察到這種情況，請提交 issue。

autocast(enabled=False) 子區域可以嵌套在啟用 autocast 的區域中。本地禁用 autocast 可能很有用，例如，如果您想強制子區域以特定的 dtype 運行。禁用 autocast 使您可以顯式控制執行類型。在子區域中，來自周圍區域的輸入應在使用前轉換為 dtype

# 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 torch.autocast(device_type="cuda"):
    e_float16 = torch.mm(a_float32, b_float32)
    with torch.autocast(device_type="cuda", enabled=False):
        # Calls e_float16.float() to ensure float32 execution
        # (necessary because e_float16 was created in an autocasted region)
        f_float32 = torch.mm(c_float32, e_float16.float())

    # No manual casts are required when re-entering the autocast-enabled region.
    # torch.mm again runs in float16 and produces float16 output, regardless of input types.
    g_float16 = torch.mm(d_float32, f_float32)

autocast 狀態是線程本地的。如果您想在新線程中啟用它，則必須在該線程中調用上下文管理器或裝飾器。 這會影響 torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel，當每個進程使用多個 GPU 時（請參閱使用多個 GPU）。

參數

• device_type (str, required) – 要使用的設備類型。可能的值為：'cuda'、'cpu'、'xpu' 和 'hpu'。該類型與 torch.device 的 type 屬性相同。 因此，您可以使用 Tensor.device.type 取得 tensor 的設備類型。

• enabled (bool, optional) – 是否應在此區域中啟用自動轉換。預設值：True

• dtype (torch_dtype, optional) – 在 autocast 中運行的 ops 的資料類型。如果 dtype 為 None，則它使用預設值（CUDA 為 torch.float16，CPU 為 torch.bfloat16），由 get_autocast_dtype() 給定。預設值：None

• cache_enabled (bool, optional) – 是否應啟用 autocast 內部的權重緩存。預設值：True

torch.amp.custom_fwd(fwd=None, *, device_type, cast_inputs=None)

為自定義 autograd 函數的 forward 方法創建一個輔助裝飾器。

Autograd 函數是 torch.autograd.Function 的子類。 有關更多詳細訊息，請參閱範例頁面。

參數

• device_type (str) – 要使用的設備類型。 'cuda'、'cpu'、'xpu' 等等。 該類型與 torch.device 的 type 屬性相同。 因此，您可以使用 Tensor.device.type 取得 tensor 的設備類型。

• cast_inputs (torch.dtype 或 None，可選，預設值=None) – 如果不是 None，當 forward 在啟用 autocast 的區域中運行時，將傳入的浮點 Tensor 轉換為目標 dtype（非浮點 Tensor 不受影響），然後在禁用 autocast 的情況下執行 forward。 如果 None，則 forward 的內部 ops 以目前的 autocast 狀態執行。

注意

如果在啟用 autocast 的區域之外調用裝飾的 forward，則 custom_fwd 是一個無操作，並且 cast_inputs 沒有效果。

torch.amp.custom_bwd(bwd=None, *, device_type)

為自定義 autograd 函數的反向傳播方法建立輔助裝飾器。

Autograd 函數是 torch.autograd.Function 的子類。 確保 backward 以與 forward 相同的 autocast 狀態執行。 有關更多詳細訊息，請參閱範例頁面。

參數

device_type (str) – 要使用的設備類型。 'cuda'、'cpu'、'xpu' 等等。 該類型與 torch.device 的 type 屬性相同。 因此，您可以使用 Tensor.device.type 取得 tensor 的設備類型。

class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)

請參閱 torch.autocast。

torch.cuda.amp.autocast(args...) 已棄用。 請改用 torch.amp.autocast("cuda", args...)。

torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None)

torch.cuda.amp.custom_fwd(args...) 已棄用。 請改用 torch.amp.custom_fwd(args..., device_type='cuda')。

torch.cuda.amp.custom_bwd(bwd)

torch.cuda.amp.custom_bwd(args...) 已棄用。 請改用 torch.amp.custom_bwd(args..., device_type='cuda')。

class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True)

請參閱 torch.autocast。 torch.cpu.amp.autocast(args...) 已棄用。 請改用 torch.amp.autocast("cpu", args...)。

梯度縮放

如果特定運算的 forward pass 具有 float16 輸入，則該運算的 backward pass 將產生 float16 梯度。 小量級的梯度值可能無法以 float16 表示。 這些值將被歸零（“下溢”），因此相應參數的更新將遺失。

為了防止下溢，“梯度縮放”將網路的損失乘以一個比例因子，並在縮放後的損失上調用 backward pass。 通過網路反向傳播的梯度然後被相同的因子縮放。 換句話說，梯度值具有更大的量級，因此它們不會歸零。

每個參數的梯度（.grad 屬性）應在優化器更新參數之前取消縮放，因此比例因子不會干擾學習率。

注意

AMP/fp16 可能不適用於每個模型！ 例如，大多數 bf16 預訓練模型無法在最大值為 65504 的 fp16 數值範圍內運行，並且會導致梯度溢位而不是下溢。 在這種情況下，比例因子可能會降至 1 以下，以嘗試將梯度帶到一個可在 fp16 動態範圍內表示的數字。 雖然人們可能期望比例始終高於 1，但我們的 GradScaler 並不保證這一點以維持效能。 如果您在使用 AMP/fp16 運行時在損失或梯度中遇到 NaN，請驗證您的模型是否相容。

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

請參閱 torch.amp.GradScaler。torch.cuda.amp.GradScaler(args...) 已棄用。請改用 torch.amp.GradScaler("cuda", args...)。

class torch.cpu.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)

請參閱 torch.amp.GradScaler。torch.cpu.amp.GradScaler(args...) 已棄用。請改用 torch.amp.GradScaler("cpu", args...)。

Autocast Op 參考

Op 資格

以 float64 或非浮點數資料類型執行的 Ops 不符合資格，無論是否啟用 autocast，都將以這些類型執行。

只有異地 (out-of-place) 的 Ops 和 Tensor 方法才符合資格。在啟用 autocast 的區域中允許使用原地 (in-place) 變體和顯式提供 out=... Tensor 的呼叫，但不會經過 autocasting。例如，在啟用 autocast 的區域中，a.addmm(b, c) 可以 autocast，但 a.addmm_(b, c) 和 a.addmm(b, c, out=d) 則不能。為了獲得最佳效能和穩定性，在啟用 autocast 的區域中，請優先使用異地的 Ops。

使用顯式的 dtype=... 參數呼叫的 Ops 不符合資格，並且將產生符合 dtype 參數的輸出。

CUDA Op 特定行為

以下列表描述了在啟用 autocast 區域中符合資格的 Ops 的行為。無論這些 Ops 是作為 torch.nn.Module 的一部分、作為函數或作為 torch.Tensor 方法調用，它們始終會經過 autocasting。如果函數在多個命名空間中公開，則無論命名空間如何，它們都會經過 autocasting。

未在下面列出的 Ops 不會經過 autocasting。它們以其輸入定義的類型執行。但是，如果未列出的 Ops 位於 autocasted Ops 的下游，autocasting 仍然可能會更改它們運行的類型。

如果某個 Op 未列出，我們假設它在 float16 中具有數值穩定性。如果您認為未列出的 Op 在 float16 中數值不穩定，請提交 issue。

可以 autocast 為 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

可以 autocast 為 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、triplet_margin_loss

CUDA 操作會提升到最寬的輸入類型

這些操作並不需要特定的 dtype 來保持穩定，但需要多個輸入，並且要求輸入的 dtype 必須匹配。如果所有輸入都是 float16，則操作會以 float16 執行。如果任何輸入是 float32，自動轉換會將所有輸入轉換為 float32，並以 float32 執行操作。

addcdiv、addcmul、atan2、bilinear、cross、dot、grid_sample、index_put、scatter_add、tensordot

此處未列出的一些操作（例如，二元操作如 add）會在本機提升輸入，而無需自動轉換的介入。如果輸入是 float16 和 float32 的混合，則無論是否啟用自動轉換，這些操作都會以 float32 執行並產生 float32 輸出。

建議使用 binary_cross_entropy_with_logits 而不是 binary_cross_entropy

torch.nn.functional.binary_cross_entropy()（以及包裝它的 torch.nn.BCELoss）的反向傳遞可能會產生無法以 float16 表示的梯度。在啟用自動轉換的區域中，正向輸入可能是 float16，這意味著反向梯度必須可以以 float16 表示（將自動轉換的 float16 正向輸入轉換為 float32 沒有幫助，因為該轉換必須在反向傳遞中還原）。因此，binary_cross_entropy 和 BCELoss 在啟用自動轉換的區域中會引發錯誤。

許多模型會在二元交叉熵層之前使用 sigmoid 層。在這種情況下，使用 torch.nn.functional.binary_cross_entropy_with_logits() 或 torch.nn.BCEWithLogitsLoss 將這兩層結合起來。binary_cross_entropy_with_logits 和 BCEWithLogits 可以安全地自動轉換型別 (autocast)。

XPU 特定運算行為（實驗性）

以下列表描述了在啟用 autocast 區域中符合資格的 Ops 的行為。無論這些 Ops 是作為 torch.nn.Module 的一部分、作為函數或作為 torch.Tensor 方法調用，它們始終會經過 autocasting。如果函數在多個命名空間中公開，則無論命名空間如何，它們都會經過 autocasting。

未在下面列出的 Ops 不會經過 autocasting。它們以其輸入定義的類型執行。但是，如果未列出的 Ops 位於 autocasted Ops 的下游，autocasting 仍然可能會更改它們運行的類型。

如果某個 Op 未列出，我們假設它在 float16 中具有數值穩定性。如果您認為未列出的 Op 在 float16 中數值不穩定，請提交 issue。

可以自動轉換型別為 float16 的 XPU 運算

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、RNNCell

可以自動轉換型別為 float32 的 XPU 運算

__pow__、__rdiv__、__rpow__、__rtruediv__、binary_cross_entropy_with_logits、cosine_embedding_loss、cosine_similarity、cumsum、dist、exp、group_norm、hinge_embedding_loss、kl_div、l1_loss、layer_norm、log、log_softmax、margin_ranking_loss、nll_loss、normalize、poisson_nll_loss、pow、reciprocal、rsqrt、soft_margin_loss、softmax、softmin、sum、triplet_margin_loss

XPU 運算會提升到最寬的輸入型別

這些操作並不需要特定的 dtype 來保持穩定，但需要多個輸入，並且要求輸入的 dtype 必須匹配。如果所有輸入都是 float16，則操作會以 float16 執行。如果任何輸入是 float32，自動轉換會將所有輸入轉換為 float32，並以 float32 執行操作。

bilinear、cross、grid_sample、index_put、scatter_add、tensordot

此處未列出的一些操作（例如，二元操作如 add）會在本機提升輸入，而無需自動轉換的介入。如果輸入是 float16 和 float32 的混合，則無論是否啟用自動轉換，這些操作都會以 float32 執行並產生 float32 輸出。

CPU 特定運算行為

以下列表描述了在啟用 autocast 區域中符合資格的 Ops 的行為。無論這些 Ops 是作為 torch.nn.Module 的一部分、作為函數或作為 torch.Tensor 方法調用，它們始終會經過 autocasting。如果函數在多個命名空間中公開，則無論命名空間如何，它們都會經過 autocasting。

未在下面列出的 Ops 不會經過 autocasting。它們以其輸入定義的類型執行。但是，如果未列出的 Ops 位於 autocasted Ops 的下游，autocasting 仍然可能會更改它們運行的類型。

如果運算未列出，我們假設它在 bfloat16 中數值上是穩定的。如果您認為未列出的運算在 bfloat16 中數值上不穩定，請提交 issue。float16 共用 bfloat16 的列表。

可以自動轉換型別為 bfloat16 的 CPU 運算

conv1d、conv2d、conv3d、bmm、mm、linalg_vecdot、baddbmm、addmm、addbmm、linear、matmul、_convolution、conv_tbc、mkldnn_rnn_layer、conv_transpose1d、conv_transpose2d、conv_transpose3d、prelu、scaled_dot_product_attention、_native_multi_head_attention

可以自動轉換型別為 float32 的 CPU 運算

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、cholesky_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、cosine_embedding_loss、nll_loss、nll_loss2d、hinge_embedding_loss、poisson_nll_loss、cross_entropy_loss、l1_loss、huber_loss、margin_ranking_loss、soft_margin_loss、triplet_margin_loss、multi_margin_loss、ctc_loss、kl_div、multilabel_margin_loss、binary_cross_entropy_with_logits、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_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、geqrf、_lu_with_info、qr、svd、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運算會提升到最寬輸入類型

這些運算不需要特定的資料類型來保持穩定性，但需要多個輸入，並且要求輸入的資料類型匹配。如果所有輸入都是 bfloat16，則運算將以 bfloat16 執行。 如果任何輸入是 float32，autocast 會將所有輸入轉換為 float32 並以 float32 執行運算。

cat、stack、index_copy

此處未列出的一些運算（例如，二元運算，如 add）本機提升輸入，而無需 autocast 的干預。 如果輸入是 bfloat16 和 float32 的混合，則這些運算以 float32 執行並產生 float32 輸出，無論是否啟用 autocast。