自动混合精度包 - torch.amp

警告

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...) 替代。

• 自动类型转换

• 梯度缩放

• 自动类型转换操作参考

  • 资格条件

  • CUDA 特定操作的行为

    • 可以自动 casts 为 float16 的 CUDA 操作

    • 可以自动 casts 为 float32 的 CUDA 操作

    • 推广到最宽输入类型的 CUDA 操作

    • 建议使用binary_cross_entropy_with_logits而非binary_cross_entropy

  • XPU 操作特定行为（实验性）

    • 可以自动转换为 float16 的 XPU 操作

    • 可自动转换为 float32 的 XPU 操作

    • 推广到最宽输入类型的 XPU 操作

  • CPU 操作的特定行为

    • 可以自动 casts 为 bfloat16 的 CPU 操作

    • 可以自动转换为 float32 的 CPU 操作

    • 提升至最宽输入类型的CPU操作

自动类型转换

torch.amp.autocast_mode.is_autocast_available(device_type)[源代码]

返回一个布尔值，指示 device_type 设备上是否支持自动类型转换（autocast）。

参数

device_type (str) – 设备类型。可能的值包括：‘cuda’, ‘cpu’, ‘xpu’ 等等。该类型与 torch.device 的 type 属性相同，因此你可以通过 Tensor.device.type 获取张量的设备类型。

返回类型

bool

classtorch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None)[源代码]

实例化的 autocast 用作上下文管理器或装饰器，使你的脚本中的一些区域可以以混合精度运行。

在这些区域中，操作会使用自动类型转换选定的特定数据类型来运行，从而提升性能并保持准确性。详细信息请参阅自动类型转换操作参考。

当进入启用了autocast的区域时，张量可以是任何类型。在使用autocasting的情况下，你不应该调用half()或bfloat16()方法来处理你的模型或输入数据。

autocast 应仅用于包裹网络的前向传递（包括损失计算），不建议在 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的区域内生成的浮点张量可能是float16类型。返回到禁用autocast的区域后，使用不同数据类型的浮点张量可能会导致类型不匹配错误。如果出现这种情况，请将autocast区域内生成的张量重新转换为float32（或所需的其他数据类型）。如果来自autocast区域的张量已经是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))

在启用了自动类型转换的区域内，如果出现了类型不匹配错误，则这是一个 bug。请在这种情况下提交一个问题报告。

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)

autocast 状态是线程本地的。如果你想在新线程中启用它，需要在线程中调用上下文管理器或装饰器。这会影响 torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel，当每个进程使用多个 GPU 时（参见与多 GPU 协同工作）。

参数

• device_type (str, required) – 设备类型。可能的值为：‘cuda’, ‘cpu’, ‘xpu’ 和 ‘hpu’。该类型与torch.device 的 type 属性相同，因此你可以通过 Tensor.device.type 获取张量的设备类型。

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

• dtype (torch_dtype, optional) – 数据类型，用于autocast中运行的操作。如果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)[源代码]

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

Autograd 函数是 torch.autograd.Function 的子类。更多详情请参见示例页面。

参数

• device_type (str) – 设备类型，如 ‘cuda’, ‘cpu’, ‘xpu’ 等。该类型与 torch.device 的 type 属性相同。因此，你可以通过 Tensor.device.type 获取张量的设备类型。

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

注意

如果装饰过的forward在未启用autocast的区域中被调用，custom_fwd将不做任何操作，并且cast_inputs也不会产生影响。

torch.amp.custom_bwd(bwd=None, *, device_type)[源代码]

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

Autograd 函数是 torch.autograd.Function 的子类。它确保 backward 方法在与 forward 相同的 autocast 状态下执行。更多详情请参阅 示例页面。

参数

device_type (str) – 设备类型，如 ‘cuda’, ‘cpu’, ‘xpu’ 等。该类型与 torch.device 的 type 属性相同。因此，你可以通过 Tensor.device.type 获取张量的设备类型。

类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') 替代。

类torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True)[源代码]

参见torch.autocast。torch.cpu.amp.autocast(args...) 已被弃用。请使用 torch.amp.autocast("cpu", args...) 代替。

梯度缩放

如果某个操作的前向传递使用了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) [源代码](<_modules/torch/cuda/amp/grad_scaler.html#GradScaler)

参见 torch.amp.GradScaler。已弃用 torch.cuda.amp.GradScaler(args...)，请使用 torch.amp.GradScaler("cuda", args...) 代替。

类 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...) 代替。

自动类型转换操作参考