自动混合精度包 - torch.amp ¶
译者:片刻小哥哥
torch.amp
提供了混合精度的便捷方法,其中某些操作使用 torch.float32
(float
) 数据类型,其他操作使用较低精度的浮点数据类型(lower_ precision_fp
):torch.float16
(half
)或torch.bfloat16
。一些操作,例如线性层和卷积,在“lower_ precision_fp”中要快得多。其他操作(例如归约)通常需要 float32
的动态范围。混合精度尝试将每个操作与其适当的数据类型相匹配。
通常,数据类型为 torch.float16
的“自动混合精度训练”会同时使用 torch.autocast
和 torch.cuda.amp.GradScaler
,如 CUDA 自动混合精度示例 和 CUDA 自动混合精度配方 中所示。 但是,torch.autocast
和 [torch.cuda.amp.GradScaler
] (#torch.cuda.amp.GradScaler "torch.cuda.amp.GradScaler") 是模块化的,如果需要,可以单独使用。 如 torch.autocast
的CPU示例部分所示,数据类型为 torch.bfloat16
的CPU上的 “自动混合精度训练/推理” 仅使用 torch.autocast
。
对于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
和 torch.cpu.amp.autocast
是1.10 版本中的新功能。
Autocasting ¶
CLASS
torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None)
[source]¶
autocast
的实例充当上下文管理器或装饰器,允许脚本区域以混合精度运行。
在这些区域中,操作以 autocast 选择的操作特定数据类型运行,以提高性能,同时保持准确性。有关详细信息,请参阅 Autocast Op 参考。
当进入启用自动转换的区域时,tensor可以是任何类型。使用自动转换时,不应在模型或输入上调用“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()
请参阅 CUDA 自动混合精度示例,了解在更复杂的场景(例如,梯度惩罚、多个模型/损失、自定义自动梯度函数)中的使用情况(以及梯度缩放)。
autocast
也可以用作装饰器,例如,在模型的 forward
方法上:
class AutocastModel(nn.Module):
...
@torch.autocast(device_type="cuda")
def forward(self, input):
...
在启用自动转换的区域中生成的浮点tensor可能是“float16”。返回到禁用自动转换的区域后,将它们与不同数据类型的浮点tensor一起使用可能会导致类型不匹配错误。如果是这样,则将自动转换区域中生成的tensor转换回“float32”(或其他需要的数据类型)。如果自动转换区域中的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(enabled=False)
子区域可以嵌套在启用了 autocast 的区域中。本地禁用 autocast 可能很有用,例如,如果您想强制子区域在特定的 dtype
中运行。禁用自动转换可以让您显式控制执行类型。在子区域中,来自周围区域的输入在使用前应转换为“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)
自动转换状态是线程本地的。如果您希望在新线程中启用它,则必须在该线程中调用上下文管理器或装饰器。这会影响 torch.nn.DataParallel
和 torch.nn.parallel.DistributedDataParallel
当每个进程与多个 GPU 一起使用时(请参阅 使用多个 GPU )。
Parameters
- device_type ( str, required* ) – 要使用的设备类型。 可能的值为:“cuda”、“cpu”、“xpu”和“hpu”。 该类型与
torch.device
的 type 属性相同。 因此,您可以使用 Tensor.device.type 获取tensor的设备类型。 - enabled ( bool, optional ) – 是否应在该区域启用自动广播。 默认值:
True
- dtype ( torch_dtype, optional ) – 是否使用 torch.float16 或 torch.bfloat16。
- cache_enabled ( bool, optional ) – 是否应启用自动投射内的权重缓存。 默认值:
True
CLASS
torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)
[source]¶
请参阅 torch.autocast
, torch.cuda.amp.autocast(args...)
相当于 torch.autocast("cuda", args...)
torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None)
[source]¶
自定义 autograd 函数的 forward
方法的辅助装饰器( torch.autograd.Function
的子类)。有关更多详细信息,请参阅示例页面。
Parameters
- cast_inputs (
torch.dtype
或 None,optional,default=None) – 如果不是None
,则当forward
运行时在启用自动转换的区域中,将传入的浮点 CUDA tensor转换为目标 dtype(非浮点tensor不受影响),然后在禁用自动转换的情况下执行forward
。如果None
,则forward
的内部操作以当前自动施放状态执行。
笔记
如果在启用自动转换的区域之外调用修饰后的 forward
,则 custom_fwd
是无操作,并且 cast_inputs
没有效果。
自定义 autograd 函数的向后方法的辅助装饰器(torch.autograd.Function
的子类)。确保 backward
以相同的方式执行自动转换状态为 forward
。有关更多详细信息,请参阅示例页面。
CLASS
torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True)
[source]¶
请参阅 torch.autocast
。 torch.cpu.amp.autocast(args...)
相当于 torch.autocast("cpu", args...)
渐变缩放 ¶
如果特定操作的前向传递具有“float16”输入,则该操作的后向传递将产生“float16”梯度。小幅度的梯度值可能无法在“float16”中表示。这些值将刷新为零(“下溢” ),因此相应参数的更新将会丢失。
为了防止下溢,“梯度缩放”将网络的损失乘以比例因子,并对缩放后的损失调用向后传递。然后通过网络向后流动的梯度按相同的因子缩放。换句话说,梯度值具有较大的幅度,因此它们不会刷新为零。
每个参数的梯度(.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)
[source]¶
返回包含比例退避因子的 Python 浮点数。
返回包含比例增长因子的 Python 浮点数。
返回包含增长区间的 Python int。
返回包含当前比例的 Python 浮点数,如果禁用缩放,则返回 1.0。
警告
get_scale()
会导致 CPU-GPU 同步。
返回一个布尔值,指示该实例是否启用。
加载缩放器状态。如果禁用此实例,则 load_state_dict()
是无操作。
Parameters
- state_dict ( dict ) – 缩放器状态。应该是从调用
state_dict()
返回的对象。
将tensor或tensor列表乘以(“缩放”)缩放因子。
返回缩放后的输出。如果未启用 GradScaler
的此实例,则返回未修改的输出。
Parameters
- outputs* ( Tensor 或 Tensors的迭代) —— 按比例输出。
Parameters
- new_scale ( float ) – 用作新的缩放退避因子。
Parameters
- new_scale ( float ) – 用作新的规模增长因素。
Parameters
- new_interval ( int ) – 用作新的增长区间。
以 dict
形式返回缩放器的状态。它包含五个条目:
"scale"
- 包含当前比例的 Python float"growth_factor"
- 包含当前增长因子的 Python float"backoff_factor"
- 包含当前退避因子的 Python float"growth_interval"
- 包含当前增长间隔的 Python int"_growth_tracker"
- 包含最近连续未跳过步骤数的 Python int。
如果未启用此实例,则返回一个空字典。
笔记
如果您希望在特定迭代后检查缩放器的状态,应在 state_dict()
之后调用 update()
。
step()
执行以下两个操作:
- 内部调用
unscale_(optimizer)
(除非明确调用unscale_()
迭代早期的“优化器”)。作为unscale_()
的一部分,检查梯度是否有 infs/NaNs.2。如果未找到 inf/NaN 梯度,则使用未缩放的梯度调用“optimizer.step()”。否则,将跳过“optimizer.step()”以避免损坏Parameters。
*args
和 **kwargs
被转发到 optimizer.step()
。
返回 optimizer.step(*args, **kwargs)
的返回值。
Parameters
- optimizer ( torch.optim.Optimizer ) – 应用梯度的优化器。
- args – 任何参数。
- kwargs – 任何关键字参数。
警告
目前不支持闭包使用。
将优化器的梯度tensor除以比例因子(“取消缩放”)。
unscale_() 是可选的,适用于需要修改或检查向后传递和 step() 之间的梯度的情况。 如果未显式调用 unscale_(),则梯度将在 step() 期间自动取消缩放。
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()
Parameters
- optimizer ( torch.optim.Optimizer ) – 拥有要取消缩放的渐变的优化器。
笔记
unscale_()
不会引起 CPU-GPU 同步。
警告
unscale_()
每个优化器每个 step()
调用,并且仅在累积了该优化器指定参数的所有梯度之后。调用 unscale_()
对于给定优化器,在每个 step()
触发运行时错误。
警告
unscale_()
可能会将稀疏梯度取消缩放,替换 .grad
属性。
更新比例因子。
如果跳过任何优化器步骤,则比例将乘以“backoff_factor”以减少它。如果连续发生“growth_interval”未跳过的迭代,则将比例乘以“growth_factor”来增加它。
传递 new_scale
手动设置新的比例值。 ( new_scale
不直接使用,它用于填充 GradScaler 的内部尺度tensor。因此,如果 new_scale
是一个tensor,以后对该tensor的就地更改不会进一步影响 GradScaler 内部使用的尺度。)
Parameters
- new_scale (float 或
torch.cuda.FloatTensor
,optional,default=None) – 新比例因子。
警告
update()
只能在迭代结束时、在 scaler.step(optimizer)
之后调用已为本次迭代使用的所有优化器调用。
警告
出于性能原因,我们不检查比例因子值以避免同步,因此比例因子不保证高于 1。如果比例低于 1 和/或您在梯度或损失中看到 NaN,则可能存在问题。例如,由于动态范围不同,bf16 预训练模型通常与 AMP/fp16 不兼容。
Autocast Op 参考 ¶
Op 资格 ¶
在“float64”或非浮点数据类型中运行的操作不符合条件,并且无论是否启用自动转换都将以这些类型运行。
只有异地操作和tensor方法才符合资格。在启用自动转换的区域中允许显式提供“out=...”tensor的就地变体和调用,但不会经过自动转换。例如,启用自动投射的区域 a.addmm(b, c)
可以自动投射,但 a.addmm_(b, c)
和 a.addmm(b, c, out=d)
不能自动投射。最好性能和稳定性,更喜欢在启用自动施放的区域中进行异地操作。
使用显式“dtype=...”参数调用的操作不符合条件,并且将产生尊重“dtype”参数的输出。
CUDA Op 特定行为 ¶
以下列表描述了启用自动转换的区域中合格操作的行为。这些操作始终会经过自动转换,无论它们是作为 torch.nn.Module
.nn.Module "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
、 triplet_margin_loss
提升到最宽输入类型的 CUDA Ops ¶
这些操作不需要特定的数据类型来保证稳定性,但需要多个输入并要求输入的数据类型匹配。如果所有输入都是 float16
,则操作在 float16
中运行。如果任何输入是 float32
,则 autocast 将所有输入转换为 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.function.binary_cross_entropy()
,它包装它)可以产生无法在float16
。在启用自动转换的区域中,前向输入可能是“float16”,这意味着后向梯度必须可以用“float16”表示(将“float16”前向输入自动转换为“float32”没有帮助,因为该转换必须在向后反转)。因此,“binary_cross_entropy”和“BCELoss”在启用自动转换的区域中会引发错误。
许多模型在二元交叉熵层之前使用 sigmoid 层。 在这种情况下,使用 torch.nn.function.binary_cross_entropy_with_logits()
或 torch.nn.BCEWithLogitsLoss
组合两个层。 binary_cross_entropy_with_logits
和 BCEWithLogits
可以安全地自动转换。
CPU Op 特定行为 ¶
以下列表描述了启用自动转换的区域中合格操作的行为。这些操作始终会经过自动转换,无论它们是作为 torch.nn.Module
.nn.Module "torch.nn.Module") ,作为函数,或作为 torch.Tensor
方法。如果函数在多个命名空间中公开,则无论命名空间如何,它们都会经历自动转换。
下面未列出的操作不会经过自动施放。它们按照输入定义的类型运行。但是,如果未列出的操作位于自动转换操作的下游,自动转换仍可能会更改其运行的类型。
如果一个操作未列出,我们假设它在“bfloat16”中数值稳定。如果您认为未列出的操作在“bfloat16”中数值不稳定,请提出问题。
可以自动转换为 bfloat16
的 CPU Ops¶
conv1d
、conv2d
、conv3d
、bmm
、mm
、baddbmm
、addmm
、addbmm
、linear
、matmul
、_convolution
可以自动转换为 float32
的 CPU Ops¶
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
、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
、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_svd_eig
、 linalg_eigh
、 linalg_lstsq
、 linalg_inv_ex
提升到最宽输入类型的 CPU Ops ¶
这些操作不需要特定的数据类型来保证稳定性,但需要多个输入并要求输入的数据类型匹配。如果所有输入都是“bfloat16”,则操作在“bfloat16”中运行。如果任何输入是 float32
,则 autocast 将所有输入转换为 float32
并运行 float32
中的操作。
cat
、 stack
、 index_copy
此处未列出的某些操作(例如,像“add”这样的二进制操作)本身会提升输入,而无需自动转换的干预。如果输入是“bfloat16”和“float32”的混合,则这些操作在“float32”中运行并产生“float32”输出,无论是否启用了自动转换。