目录

自动混合精度包 - torch.amp

torch.amp 提供了混合精度的便捷方法, 其中一些操作使用 torch.float32 (float) 数据类型,而其他操作 使用较低精度的浮点数数据类型 (lower_precision_fp): torch.float16 (half) 或 torch.bfloat16。某些操作,如线性层和卷积, 在 lower_precision_fp 中运行得更快。其他操作,如归约操作, 通常需要 float32 的动态范围。 混合精度尝试将每个操作与其适当的数据类型匹配。

通常,“自动混合精度训练”使用数据类型torch.float16时会结合使用torch.autocasttorch.cuda.amp.GradScaler,如在CUDA 自动混合精度示例CUDA 自动混合精度示例中所示。 然而,torch.autocasttorch.cuda.amp.GradScaler是模块化的,可以根据需要单独使用。 如在torch.autocast的CPU示例部分所示,“自动混合精度训练/推理”在CPU上使用 数据类型torch.bfloat16时仅使用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.autocasttorch.cpu.amp.autocast 是版本 1.10 中的新功能。

自动类型转换

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

autocast 的实例用作上下文管理器或装饰器,允许脚本中的某些区域以混合精度运行。

在这些区域中,ops 以自动转换选择的特定操作 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()

请参阅CUDA 自动混合精度示例以了解在更复杂场景中的使用(包括梯度缩放), 例如,梯度惩罚、多个模型/损失函数、自定义 autograd 函数。

autocast 也可以用作装饰器,例如,在你的模型的 forward 方法上:

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

在启用了自动混合精度的区域中生成的浮点张量可能是 float16。 返回到禁用自动混合精度的区域后,与不同数据类型的浮点张量一起使用可能会导致类型不匹配错误。如果是这样,请将自动混合精度区域内生成的张量重新转换为 float32(或其他所需的数据类型)。 如果来自自动混合精度区域的张量已经是 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)

CPU 推理示例与 Jit 追踪:

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))

类型不匹配错误是在自动混合精度区域中的一个 bug;如果您观察到这种情况,请提交一个问题。

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.DataParalleltorch.nn.parallel.DistributedDataParallel,当与每个进程多个GPU一起使用时(参见使用多个GPU)。

Parameters
  • 设备类型 (str, 必需) – 使用的设备类型。可能的值为:‘cuda’,‘cpu’,‘xpu’ 和 ‘hpu’。 该类型与type属性相同torch.device。 因此,您可以使用Tensor.device.type获取张量的设备类型。

  • 启用 (bool, 可选) – 是否应在该区域启用自动类型转换。 默认值:True

  • dtype (torch_dtype, 可选) – 是否使用 torch.float16 或 torch.bfloat16。

  • 启用缓存 (bool, 可选) – 是否启用autocast中的权重缓存。 默认值: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]

为自定义自动微分函数的forward方法创建一个辅助装饰器。

自动求导函数是 torch.autograd.Function 的子类。 详情请参见 示例页面

Parameters

cast_inputs (torch.dtype 或 None,可选,默认值=None) – 如果不是 None, 当 forward 在启用了自动类型转换的区域运行时,会将传入的浮点型 CUDA 张量转换为目标数据类型(非浮点型张量不受影响), 然后在禁用自动类型转换的情况下执行 forward。 如果为 Noneforward 的内部操作将使用当前的自动类型转换状态执行。

注意

如果装饰的 forward 在未启用autocast的区域中被调用, custom_fwd 将不起作用,并且 cast_inputs 没有任何效果。

torch.cuda.amp.custom_bwd(bwd)[source]

为自定义自动微分函数的反向方法创建一个辅助装饰器。

自动求导函数是 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]

一个实例 scaler of GradScaler.

方便地执行梯度缩放的步骤。

  • scaler.scale(loss) 将给定的损失乘以 scaler 的当前缩放因子。

  • scaler.step(optimizer) 安全地取消梯度缩放并调用 optimizer.step()

  • scaler.update() 更新 scaler 的缩放因子。

Example:

# 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 的值可能会“溢出”(变为 inf 或 NaN)。因此,最佳的缩放因子是在不产生 inf 或 NaN 梯度值的情况下可以使用的最大因子。 scaler 通过在每次 scaler.step(optimizer)(或可选的独立 scaler.unscale_(optimizer),参见 unscale_())期间检查梯度中的 inf 和 NaN 来随时间近似最佳缩放因子。

  • 如果检测到infs/NaNs,scaler.step(optimizer) 会跳过底层的 optimizer.step()(因此参数本身不会被破坏),并且 update() 会将比例乘以 backoff_factor

  • 如果没有发现无穷大值/非数值,scaler.step(optimizer) 将像平常一样运行底层的 optimizer.step()。 如果连续发生 growth_interval 次未跳过的迭代,update() 会将比例乘以 growth_factor

缩放因子在初始几次迭代中,由于其值的校准,常常会导致梯度中出现infs/NaNs。 scaler.step 会跳过这些迭代中的 optimizer.step() 。之后,步长跳过应很少发生(每隔几百或几千次迭代一次)。

Parameters
  • init_scale (float, 可选, 默认=2.**16) – 初始缩放因子。

  • 增长因子 (float, 可选, 默认=2.0) – 如果连续growth_interval次迭代中没有出现inf/NaN梯度,则在update()期间将缩放乘以此因子。

  • backoff_factor (float, 可选, 默认=0.5) – 如果在迭代中出现inf/NaN梯度,则乘以此因子。

  • growth_interval (int, 可选, 默认=2000) – 在必须发生连续的无 inf/NaN 梯度的迭代次数,以使比例乘以 growth_factor

  • 已启用 (布尔值, 可选) – 如果 False,禁用梯度缩放。 step() 仅调用底层的 optimizer.step(),其他方法变为无效操作。 默认: True

get_backoff_factor()[source]

返回一个包含缩放回退因子的 Python 浮点数。

Return type

浮点型

get_growth_factor()[source]

返回一个包含缩放增长因子的 Python 浮点数。

Return type

浮点型

get_growth_interval()[source]

返回一个包含增长间隔的 Python 整数。

Return type

整数

get_scale()[source]

返回一个包含当前缩放比例的 Python 浮点数,如果禁用缩放则返回 1.0。

警告

get_scale() 会引发 CPU-GPU 同步。

Return type

浮点型

is_enabled()[source]

返回一个布尔值,指示此实例是否已启用。

Return type

布尔

load_state_dict(state_dict)[source]

加载缩放器状态。

如果此实例被禁用,load_state_dict() 将不执行任何操作。

Parameters

state_dict (dict) – 标量状态。应该是一个从对 state_dict() 的调用返回的对象。

scale(outputs: Tensor) Tensor[source]
scale(outputs: List[Tensor]) List[Tensor]
scale(outputs: Tuple[Tensor, ...]) Tuple[Tensor, ...]
scale(outputs: Iterable[Tensor]) Iterable[Tensor]

将张量或张量列表乘以(“缩放”)一个缩放因子。

返回缩放后的输出。如果此实例的 GradScaler 未启用,则输出将保持不变。

Parameters

输出 (张量张量的可迭代对象张量) – 需要缩放的输出。

set_backoff_factor(new_factor)[source]

设置一个新的缩放回退因子。

Parameters

new_scale (float) – 用作新缩放回退因子的值。

set_growth_factor(new_factor)[source]

设置一个新的比例增长因子。

Parameters

new_scale (float) – 用作新尺度增长因子的值。

set_growth_interval(new_interval)[source]

设置一个新的增长间隔。

Parameters

new_interval (int) – 用作新增长间隔的值。

state_dict()[source]

返回缩放器的状态为 dict

它包含五个条目:

  • "scale" - 一个包含当前比例的 Python 浮点数

  • "growth_factor" - 一个包含当前增长因子的Python浮点数

  • "backoff_factor" - 一个包含当前回退因子的Python浮点数

  • "growth_interval" - 一个包含当前增长间隔的Python整数

  • "_growth_tracker" - 一个Python整数,表示最近连续未跳过的步骤数量。

如果此实例未启用,则返回一个空字典。

注意

如果你希望在特定迭代后保存缩放器的状态,state_dict() 应该在 update() 之后调用。

Return type

字典[字符串, 任意类型]

step(optimizer, *args, **kwargs)[source]

如果梯度不是无穷大或NaN,则调用 unscale_(optimizer) 并随后更新参数。

step() 执行以下两个操作:

  1. 内部调用 unscale_(optimizer)(除非在迭代的早期已经为 optimizer 明确调用了 unscale_())。作为 unscale_() 的一部分,会检查梯度是否存在 inf/NaN。

  2. 如果没有发现无穷大/NaN梯度,则使用未缩放的梯度调用optimizer.step()。否则,将跳过optimizer.step()以避免损坏参数。

*args**kwargs 被转发到 optimizer.step()

返回 optimizer.step(*args, **kwargs) 的返回值。

Parameters
  • 优化器 (torch.optim.Optimizer) – 应用梯度的优化器。

  • args (Any) – 任何参数。

  • kwargs (Any) – 任意关键字参数。

Return type

可选[浮点数]

警告

目前不支持闭包的使用。

unscale_(optimizer)[source]

将优化器的梯度张量除以缩放因子(“反缩放”)。

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

优化器 (torch.optim.Optimizer) – 拥有需要取消缩放梯度的优化器。

注意

unscale_() 不会导致 CPU-GPU 同步。

警告

unscale_() 应该在每个优化器的每次 step() 调用中仅调用一次, 并且仅在为该优化器分配的所有参数都累积了梯度之后。 对于给定的优化器,在每次 step() 之间两次调用 unscale_() 将触发 RuntimeError。

警告

unscale_() 可能会在原地取消稀疏梯度的缩放,替换 .grad 属性。

update(new_scale=None)[source]

更新比例因子。

如果任何优化器步骤被跳过,比例将乘以 backoff_factor 以减小它。如果连续发生 growth_interval 次未跳过的迭代, 比例将乘以 growth_factor 以增加它。

Passing new_scale sets the new scale value manually. (new_scale is not used directly, it’s used to fill GradScaler’s internal scale tensor. So if new_scale was a tensor, later in-place changes to that tensor will not further affect the scale GradScaler uses internally.)

Parameters

new_scale (float 或 torch.cuda.FloatTensor, 可选, default=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中数值不稳定, 请提交一个问题。

CUDA 运算符可以自动转换为 float16

__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

CUDA 运算符可以自动转换为 float32

__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 操作

这些操作不需要特定的数据类型来保证稳定性,但需要多个输入且输入的数据类型必须匹配。如果所有的输入都是 float16,该操作将在 float16 中运行。如果任一输入是 float32, autocast 将把所有输入转换为 float32 并在 float32 中运行操作。

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

有些操作在这里没有列出(例如,二元操作如add)会原生地提升输入,而无需自动类型转换的干预。如果输入是float16float32的混合,这些操作会在float32中运行,并产生float32输出,无论是否启用了自动类型转换。

建议使用 binary_cross_entropy_with_logits 而不是 binary_cross_entropy

反向传播过程中的第 torch.nn.functional.binary_cross_entropy() (以及包裹它的第 torch.nn.BCELoss) 可能会产生无法在 float16 中表示的梯度。在自动混合精度启用的区域中,前向输入可能是 float16,这意味着反向传播的梯度必须能够在 float16 中表示(将前向输入自动转换为 float16 并不解决问题,因为在反向传播中需要将其还原)。 因此,在自动混合精度启用的区域中,binary_cross_entropyBCELoss 会引发错误。

许多模型在二元交叉熵层之前使用一个 sigmoid 层。 在这种情况下,使用torch.nn.functional.binary_cross_entropy_with_logits()torch.nn.BCEWithLogitsLoss结合这两个层。binary_cross_entropy_with_logitsBCEWithLogits 可以安全地自动转换。

CPU 特定操作行为

以下列表描述了在自动类型转换启用区域中符合条件的操作的行为。 这些操作始终会经历自动类型转换,无论它们是作为torch.nn.Module的一部分被调用, 作为一个函数,还是作为一个torch.Tensor方法。如果函数在多个命名空间中暴露, 它们都会经历自动类型转换,而与命名空间无关。

未在下面列出的操作不会经过自动类型转换。它们按照其输入定义的类型运行。然而,如果未列出的操作位于自动类型转换操作的下游,自动类型转换仍可能会更改这些操作的运行类型。

如果一个操作未列出,我们假设它在bfloat16中数值稳定。 如果你认为一个未列出的操作在bfloat16中数值不稳定, 请提交一个问题。

可以在自动类型转换中使用的CPU运算符bfloat16

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

可以在自动类型转换中使用的CPU运算符float32

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 操作

这些操作不需要特定的数据类型来保证稳定性,但需要多个输入且输入的数据类型必须匹配。如果所有的输入都是 bfloat16,该操作将在 bfloat16 中运行。如果任一输入是 float32, autocast 将把所有输入转换为 float32 并在 float32 中运行操作。

cat, stack, index_copy

有些操作在这里没有列出(例如,二元操作如add)会原生地提升输入,而无需自动类型转换的干预。如果输入是bfloat16float32的混合,这些操作会在float32中运行,并产生float32输出,无论是否启用了自动类型转换。

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源