# 自动混合精度包 - torch.amp > 原文:[`pytorch.org/docs/stable/amp.html`](https://pytorch.org/docs/stable/amp.html)> > > 译者:[飞龙](https://github.com/wizardforcel) > > 协议:[CC BY-NC-SA 4.0](http://creativecommons.org/licenses/by-nc-sa/4.0/) `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 自动混合精度配方](https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html) 中所示。但是,如果需要,`torch.autocast` 和 `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 版本中的新功能。 + 自动转换 + 梯度缩放 + Autocast Op 参考 + Op Eligibility + CUDA Op-Specific Behavior + CUDA Ops that can autocast to `float16` + CUDA Ops that can autocast to `float32` + CUDA Ops that promote to the widest input type + 推荐使用 `binary_cross_entropy_with_logits` 而不是 `binary_cross_entropy` + CPU Op-Specific Behavior + CPU Ops that can autocast to `bfloat16` + CPU Ops that can autocast to `float32` + CPU Ops that promote to the widest input type ## 自动转换 ```py class torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None) ``` `autocast` 的实例充当上下文管理器或装饰器,允许脚本的区域以混合精度运行。 在这些区域中,操作以由 autocast 选择的 op-specific dtype 运行,以提高性能同时保持准确性。详细信息请参阅 Autocast Op 参考。 进入启用 autocast 的区域时,张量可以是任何类型。在使用自动转换时,不应在模型或输入上调用 `half()` 或 `bfloat16()`。 `autocast` 应该仅包装网络的前向传递,包括损失计算。不建议在 autocast 下进行反向传递。反向操作以 autocast 用于相应前向操作的相同类型运行。 CUDA 设备示例: ```py # 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`方法: ```py class AutocastModel(nn.Module): ... @torch.autocast(device_type="cuda") def forward(self, input): ... ``` 在自动转换启用的区域中生成的浮点张量可能是`float16`。在返回到禁用自动转换的区域后,将其与不同 dtype 的浮点张量一起使用可能会导致类型不匹配错误。如果是这样,请将在自动转换区域中生成的张量重新转换为`float32`(或其他所需的 dtype)。如果来自自动转换区域的张量已经是`float32`,则转换是一个空操作,并且不会产生额外的开销。CUDA 示例: ```py # 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 训练示例: ```py # 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 推理示例: ```py # 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 推理示例: ```py 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)`子区域可以嵌套在自动转换启用的区域中。局部禁用自动转换可能很有用,例如,如果您想要强制子区域以特定的`dtype`运行。禁用自动转换可以让您明确控制执行类型。在子区域中,应将来自周围区域的输入转换为`dtype`后再使用: ```py # 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) ``` 自动转换状态是线程本地的。如果要在新线程中启用它,则必须在该线程中调用上下文管理器或装饰器。这会影响当与一个进程中的多个 GPU 一起使用时的`torch.nn.DataParallel`和`torch.nn.parallel.DistributedDataParallel`(请参阅使用多个 GPU)。 参数 + **device_type**(*str*,*必需*)- 要使用的设备类型。可能的值为:'cuda'、'cpu'、'xpu'和'hpu'。该类型与`torch.device`的类型属性相同。因此,您可以使用 Tensor.device.type 获取张量的设备类型。 + **enabled**(*bool*,*可选*)- 是否应在该区域启用自动转换。默认值:`True` + **dtype**(*torch_dtype**,* *可选*)- 是否使用 torch.float16 或 torch.bfloat16。 + **cache_enabled**(*bool*,*可选*)- 是否启用自动转换内部的权重缓存。默认值:`True` ```py class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True) ``` 请参阅`torch.autocast`。 `torch.cuda.amp.autocast(args...)`等同于`torch.autocast("cuda", args...)` ```py torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None) ``` 为自定义自动求导函数的`forward`方法创建一个辅助装饰器。 自动求导函数是`torch.autograd.Function`的子类。有关更多详细信息,请参阅示例页面。 参数 **cast_inputs**(`torch.dtype`或 None,可选,默认为 None)- 如果不是`None`,当`forward`在自动转换启用的区域运行时,将传入的浮点 CUDA 张量转换为目标 dtype(非浮点张量不受影响),然后在禁用自动转换的情况下执行`forward`。如果为`None`,`forward`的内部操作将以当前自动转换状态执行。 注意 如果装饰的`forward`在自动转换启用的区域之外被调用,则`custom_fwd`是一个空操作,`cast_inputs`没有效果。 ```py torch.cuda.amp.custom_bwd(bwd) ``` 为自定义自动求导函数的反向方法创建一个辅助装饰器。 自动求导函数是`torch.autograd.Function`的子类。确保`backward`以与`forward`相同的自动转换状态执行。更多细节请参见示例页面。 ```py class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True) ``` 请参见`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,请验证您的模型是否兼容。 ```py class torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True) ``` `GradScaler`的一个实例`scaler`。 帮助方便地执行梯度缩放的步骤。 + `scaler.scale(loss)`会将给定的损失乘以`scaler`当前的比例因子。 + `scaler.step(optimizer)`会安全地取消梯度缩放并调用`optimizer.step()`。 + `scaler.update()`会更新`scaler`的比例因子。 示例: ```py # Creates a GradScaler once at the beginning of training. scaler = GradScaler() for epoch in epochs: for input, target in data: optimizer.zero_grad() output = model(input) loss = loss_fn(output, target) # Scales loss. Calls backward() on scaled loss to create scaled gradients. scaler.scale(loss).backward() # scaler.step() first unscales gradients of the optimizer's params. # If gradients don't contain infs/NaNs, optimizer.step() is then called, # otherwise, optimizer.step() is skipped. scaler.step(optimizer) # Updates the scale for next iteration. scaler.update() ``` 查看自动混合精度示例以获取更复杂情况下(包括自动转换)的用法,如梯度裁剪、梯度累积、梯度惩罚和多个损失/优化器。 `scaler`会动态估计每次迭代的比例因子。为了最小化梯度下溢,应该使用一个较大的比例因子。但是,如果比例因子太大,`float16`值可能会“溢出”(变为无穷大或 NaN)。因此,最佳比例因子是可以在不产生无穷大或 NaN 梯度值的情况下使用的最大因子。`scaler`通过在每次`scaler.step(optimizer)`(或可选的单独`scaler.unscale_(optimizer)`,参见`unscale_()`)期间检查梯度中的无穷大和 NaN 来随时间近似最佳比例因子。 + 如果发现无穷大/NaN,`scaler.step(optimizer)`会跳过底层的`optimizer.step()`(因此参数本身保持不受损坏),`update()`会将比例乘以`backoff_factor`。 + 如果没有发现无穷大/NaN,`scaler.step(optimizer)`会像往常一样运行底层的`optimizer.step()`。如果连续发生`growth_interval`个未跳过的迭代,`update()`会将比例乘以`growth_factor`。 比例因子通常会导致在前几次迭代中梯度中出现无穷大/NaN,因为其值在校准。对于这些迭代,`scaler.step`会跳过底层的`optimizer.step()`。之后,跳过步骤应该很少发生(每几百或几千次迭代一次)。 参数 + **init_scale**(*float*,*可选*,*默认=2.**16*)- 初始比例因子。 + **growth_factor** ([*float*](https://docs.python.org/3/library/functions.html#float "(在 Python v3.12 中)")*,* *可选**,* *默认=2.0*) – 如果在`growth_interval`连续迭代中没有 inf/NaN 梯度,则在`update()`期间将缩放乘以的因子。 + **backoff_factor** ([*float*](https://docs.python.org/3/library/functions.html#float "(在 Python v3.12 中)")*,* *可选**,* *默认=0.5*) – 如果在迭代中出现 inf/NaN 梯度,则在`update()`期间将缩放乘以的因子。 + **growth_interval** ([*int*](https://docs.python.org/3/library/functions.html#int "(在 Python v3.12 中)")*,* *可选**,* *默认=2000*) – 必须在没有 inf/NaN 梯度的连续迭代中发生的次数,以便将比例乘以`growth_factor`。 + **enabled** ([*bool*](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)")*,* *可选*) – 如果为`False`,则禁用梯度缩放。`step()`只会调用底层的`optimizer.step()`,其他方法将变为无操作。默认值:`True` ```py get_backoff_factor() ``` 返回一个包含缩放退避因子的 Python 浮点数。 返回类型 [float](https://docs.python.org/3/library/functions.html#float "(在 Python v3.12 中)") ```py get_growth_factor() ``` 返回一个包含增长因子的 Python 浮点数。 返回类型 [float](https://docs.python.org/3/library/functions.html#float "(在 Python v3.12 中)") ```py get_growth_interval() ``` 返回一个包含增长间隔的 Python 整数。 返回类型 [int](https://docs.python.org/3/library/functions.html#int "(在 Python v3.12 中)") ```py get_scale() ``` 返回一个包含当前比例的 Python 浮点数,如果禁用了缩放,则返回 1.0。 警告 `get_scale()` 会导致 CPU-GPU 同步。 返回类型 [float](https://docs.python.org/3/library/functions.html#float "(在 Python v3.12 中)") ```py is_enabled() ``` 返回一个指示此实例是否已启用的布尔值。 返回类型 [bool](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)") ```py load_state_dict(state_dict) ``` 加载缩放器状态。 如果此实例已禁用,则`load_state_dict()`将不执行任何操作。 参数 **state_dict** ([*dict*](https://docs.python.org/3/library/stdtypes.html#dict "(在 Python v3.12 中)")) – 缩放器状态。应该是从调用`state_dict()`返回的对象。 ```py scale(outputs: Tensor) → Tensor ``` ```py scale(outputs: List[Tensor]) → List[Tensor] ``` ```py scale(outputs: Tuple[Tensor, ...]) → Tuple[Tensor, ...] ``` ```py scale(outputs: Iterable[Tensor]) → Iterable[Tensor] ``` 通过比例因子将张量或张量列表进行乘法(“缩放”)。 返回缩放后的输出。如果此`GradScaler`实例未启用,则输出将保持不变。 参数 **outputs** (*Tensor* *或* *可迭代* *的* *Tensors*) – 要缩放的输出。 ```py set_backoff_factor(new_factor) ``` 设置一个新的退避因子。 参数 **new_scale** ([*float*](https://docs.python.org/3/library/functions.html#float "(在 Python v3.12 中)")) – 用作新的缩放退避因子的值。 ```py set_growth_factor(new_factor) ``` 设置一个新的增长因子。 参数 **new_scale** ([*float*](https://docs.python.org/3/library/functions.html#float "(在 Python v3.12 中)")) – 用作新的增长因子的值。 ```py set_growth_interval(new_interval) ``` 设置一个新的增长间隔。 参数 **new_interval** ([*int*](https://docs.python.org/3/library/functions.html#int "(在 Python v3.12 中)")) – 用作新增长间隔的值。 ```py state_dict() ``` 返回一个[`dict`](https://docs.python.org/3/library/stdtypes.html#dict "(在 Python v3.12 中)")形式的缩放器状态。 它包含五个条目: + `"scale"` - 包含当前比例的 Python 浮点数 + `"growth_factor"` - 包含当前增长因子的 Python 浮点数 + `"backoff_factor"` - 包含当前退避因子的 Python 浮点数 + `"growth_interval"` - 包含当前增长间隔的 Python 整数 + `"_growth_tracker"` - 包含最近连续未跳过步骤的 Python int。 如果此实例未启用,则返回空字典。 注意 如果您希望在特定迭代之后检查标量器的状态,应在`update()`之后调用`state_dict()`。 返回类型 [*Dict*](https://docs.python.org/3/library/typing.html#typing.Dict "(in Python v3.12)")[[str](https://docs.python.org/3/library/stdtypes.html#str "(in Python v3.12)"), [*Any*](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.12)")] ```py step(optimizer, *args, **kwargs) ``` 如果梯度不是无穷大/NaN,则调用`unscale_(optimizer)`,然后进行参数更新。 `step()`执行以下两个操作: 1. 在迭代中未为`optimizer`显式调用`unscale_()`时,内部调用`unscale_(optimizer)`。作为`unscale_()`的一部分,梯度将被检查是否有无穷大/NaN。 1. 如果未发现无穷大/NaN 梯度,则使用未缩放梯度调用`optimizer.step()`。否则,将跳过`optimizer.step()`以避免损坏参数。 `*args`和`**kwargs`将被转发到`optimizer.step()`。 返回`optimizer.step(*args, **kwargs)`的返回值。 参数 + **optimizer** (*torch.optim.Optimizer*) – 应用梯度的优化器。 + **args** ([*Any*](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.12)")) – 任意参数。 + **kwargs** ([*Any*](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.12)")) – 任意关键字参数。 返回类型 [*Optional*](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.12)")[[float](https://docs.python.org/3/library/functions.html#float "(in Python v3.12)")] 警告 目前不支持闭包使用。 ```py unscale_(optimizer) ``` 通过比例因子除(“取消缩放”)优化器的梯度张量。 `unscale_()`是可选的,用于需要在反向传播和`step()`之间修改或检查梯度的情况。如果未显式调用`unscale_()`,梯度将在`step()`期间自动取消缩放。 简单示例,使用`unscale_()`启用未缩放梯度的裁剪: ```py ... 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 同步。 警告 每次调用`step()`时,应仅对每个优化器调用一次`unscale_()`,并且仅在为该优化器分配的参数累积了所有梯度之后才调用。在每个`step()`之间两次为给定优化器调用`unscale_()`会触发运行时错误。 警告 `unscale_()`可能会在原地取消缩放稀疏梯度,替换`.grad`属性。 ```py update(new_scale=None) ``` 更新比例因子。 如果有任何优化器步骤被跳过,则将比例乘以 `backoff_factor` 以减小它。如果连续发生 `growth_interval` 个未跳过的迭代,则将比例乘以 `growth_factor` 以增加它。 传递 `new_scale` 可以手动设置新的比例值。(`new_scale` 不会直接使用,它用于填充 GradScaler 的内部比例张量。因此,如果 `new_scale` 是一个张量,则稍后对该张量的就地更改不会进一步影响 GradScaler 内部使用的比例。) 参数 **new_scale** (float 或 `torch.cuda.FloatTensor`, 可选, 默认=None) – 新的比例因子。 警告 `update()` 应该只在迭代结束后调用,在此迭代中已为所有使用的优化器调用了 `scaler.step(optimizer)`。 警告 出于性能原因,我们不检查比例因子的值以避免同步,因此不能保证比例因子大于 1。如果比例低于 1 和/或在梯度或损失中看到 NaN,则可能有问题。例如,由于动态范围不同,bf16 预训练模型通常与 AMP/fp16 不兼容。 ## 自动混合精度操作参考 ### 操作资格 在 `float64` 或非浮点数据类型中运行的操作不符合条件,无论是否启用了自动混合精度,它们都将以这些类型运行。 只有不改变原始数据的操作和张量方法才符合条件。在启用自动混合精度的区域中,可以使用就地变体和显式提供 `out=...` 张量的调用,但不会经过自动混合精度。例如,在启用自动混合精度的区域中,`a.addmm(b, c)` 可以自动混合精度,但 `a.addmm_(b, c)` 和 `a.addmm(b, c, out=d)` 不能。为了获得最佳性能和稳定性,请在启用自动混合精度的区域中使用不改变原始数据的操作。 使用显式 `dtype=...` 参数调用的操作不符合条件,并将生成符合 `dtype` 参数的输出。 ### CUDA 操作特定行为 以下列表描述了在启用自动混合精度的区域中符合条件的操作的行为。这些操作总是经过自动混合精度,无论它们作为 `torch.nn.Module` 的一部分调用,作为一个函数调用,还是作为 `torch.Tensor` 方法调用。如果函数在多个命名空间中公开,无论命名空间如何,它们都会经过自动混合精度。 下面未列出的操作不会经过自动混合精度。它们以其输入定义的类型运行。但是,如果它们在自动混合精度操作的下游,则自动混合精度可能会改变未列出操作运行的类型。 如果某个操作未列出,我们会假定它在 `float16` 中是数值稳定的。如果您认为某个未列出的操作在 `float16` 中是数值不稳定的,请提交一个问题。 #### 可以自动混合精度为 `float16` 的 CUDA 操作 `__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 操作 `__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`可以安全地进行自动转换。 ### CPU 操作特定行为 以下列表描述了在启用自动转换的区域中合格操作的行为。这些操作始终经过自动转换,无论它们是作为`torch.nn.Module`的一部分、作为函数还是作为`torch.Tensor`方法调用。如果函数在多个命名空间中公开,无论命名空间如何,它们都会经过自动转换。 以下未列出的操作不会经过自动转换。它们以其输入定义的类型运行。但是,如果它们在自动转换的操作之后,自动转换仍可能更改未列出操作运行的类型。 如果一个操作未列出,我们假定它在`bfloat16`中是数值稳定的。如果您认为未列出的操作在`bfloat16`中是数值不稳定的,请提交一个问题。 #### CPU 操作可以自动转换为`bfloat16` `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`,`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_eig`,`linalg_eigh`,`linalg_lstsq`,`linalg_inv_ex` #### 将输入类型提升为最宽的 CPU 操作 这些操作不需要特定的 dtype 来保持稳定性,但需要多个输入并要求输入的 dtypes 匹配。如果所有输入都是`bfloat16`,则该操作将在`bfloat16`中运行。如果任何输入是`float32`,自动转换将将所有输入转换为`float32`并在`float32`中运行该操作。 `cat`,`stack`,`index_copy` 这里未列出一些操作(例如,像`add`这样的二进制操作)会在没有自动转换干预的情况下本地提升输入。如果输入是`bfloat16`和`float32`的混合,这些操作将在`float32`中运行并产生`float32`输出,无论是否启用了自动转换。