Pytorch-API：混合精度计算（AMP）

混合精度计算

Automatic Mixed Precision package - torch.amp
torch.amp 提供了混合精度的便捷方法，其中一些操作使用 torch.float32（浮点）数据类型，而其他操作使用较低精度的浮点数据类型（lower_precision_fp）：torch.float16（半精度）或 torch.bfloat16。某些操作（如线性层和卷积层）在 lower_precision_fp 下速度更快，而其他操作（如归约）通常需要 float32 的动态范围。混合精度尝试将每个操作匹配到其适合的数据类型。
通常，“自动混合精度训练”使用数据类型 torch.float16 时，会结合使用 torch.autocast 和 torch.cpu.amp.GradScaler 或 torch.cuda.amp.GradScaler，如 CUDA 自动混合精度示例和 CUDA 自动混合精度指南中所示。然而，torch.autocast 和 torch.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.GradScaler(“cuda”, args…) 等同于 torch.cuda.amp.GradScaler(args…)。

• torch.GradScaler(“cpu”, args…) 等同于 torch.cpu.amp.GradScaler(args…)。

torch.autocast 和 torch.cpu.amp.autocast 是在版本 1.10 中引入的新功能。

1.Autocasting

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

Parameters:
	device_type (str, required) – Device type to use. Possible values are: ‘cuda’, ‘cpu’, ‘xpu’ and ‘hpu’. The type is the same as the type attribute of a torch.device. Thus, you may obtain the device type of a tensor using Tensor.device.type.
	
	enabled (bool, optional) – Whether autocasting should be enabled in the region. Default: True
	
	dtype (torch_dtype, optional) – Whether to use torch.float16 or torch.bfloat16.
	
	cache_enabled (bool, optional) – Whether the weight cache inside autocast should be enabled. Default: True

• autocast 实例作为上下文管理器或装饰器，允许脚本中的某些区域以混合精度运行。
• 在这些区域中，操作会以 autocast 选择的特定数据类型运行，以提高性能，同时保持精度。详情请参见 Autocast 操作参考。
• 进入 autocast 启用区域时，张量可以是任何类型。使用自动混合精度时，不应在模型或输入上调用 half() 或 bfloat16()。
• autocast 应该只包裹网络的前向传递（包括损失计算）。不建议在 autocast 下运行反向传递。反向操作将使用 autocast 为对应的前向操作使用的相同数据类型。

Example for CUDA Devices:

# 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自动混合精度计算示例以了解在更复杂场景中的使用（例如，梯度惩罚、多个模型/损失、自定义自动求导函数）及其与梯度缩放的结合。

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

在 autocast 启用区域中生成的浮点张量可能是 float16。返回到 autocast 禁用区域后，将它们与不同数据类型的浮点张量一起使用可能会导致类型不匹配错误。在这种情况下，请将 autocast 区域中生成的张量转换回 float32（或其他所需的数据类型）。如果 autocast 区域中的张量已经是 float32，那么转换操作是无操作（no-op），不会产生额外的开销。

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 training 示例

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

在启用了 autocast 的区域中出现类型不匹配错误是一个 bug;
在启用了 autocast 的区域中可以嵌套 autocast(enabled=False) 子区域。局部禁用 autocast 可以很有用，例如，如果您想要强制子区域以特定的数据类型运行。禁用 autocast 可以让您对执行类型有明确的控制。在子区域中，应在使用之前将来自周围区域的输入转换为相应的数据类型：

# 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》）。

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

• torch.cuda.amp.autocast(args…) is equivalent to torch.autocast(“cuda”, args…)

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

-参数
-cast_inputs（torch.dtype 或 None，可选，默认为 None） - 如果不为 None，则在 forward 在自动混合精度启用区域运行时，将传入的浮点 CUDA 张量转换为目标数据类型（非浮点张量不受影响），然后在禁用 autocast 的情况下执行 forward。如果为 None，则 forward 的内部操作将根据当前 autocast 状态执行。
创建一个用于自定义自动求导函数的前向方法的辅助装饰器。

自动求导函数是 torch.autograd.Function 的子类。更多详细信息请参见示例。
如果在没有启用自动混合精度的区域中调用被装饰的 forward 方法，则 custom_fwd 不会执行任何操作，并且 cast_inputs 不会起作用。

1.4 torch.cuda.amp.custom_bwd(bwd)

创建一个用于自定义自动求导函数的后向方法的辅助装饰器。

自动求导函数是 torch.autograd.Function 的子类。确保后向方法的执行与前向方法的自动混合精度状态相同。更多详细信息请参见示例。

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

torch.cpu.amp.autocast(args…) is equivalent to torch.autocast(“cpu”, args…)

2 梯度缩放

如果某个操作的前向传递具有 float16 输入，则该操作的后向传递将产生 float16 梯度。具有较小幅度的梯度值可能无法表示为 float16。这些值将被截断为零（“欠流动”），因此相应参数的更新将丢失。

为防止欠流动，“梯度缩放”通过将网络的损失（们）乘以一个比例因子，并对缩放后的损失（们）执行反向传递来解决。通过网络向后传播的梯度然后被同样的因子缩放。换句话说，梯度值具有较大的幅度，因此它们不会被截断为零。

每个参数的梯度（.grad 属性）在优化器更新参数之前应该被还原，以便比例因子不会干扰学习率。

AMP/fp16 可能并不适用于每个模型！例如，大多数 bf16 预训练模型无法在 fp16 数值范围的最大值 65504 内运行，而是会导致梯度溢出而不是欠流动。在这种情况下，缩放因子可能会降低到 1 以下，以试图将梯度带到 fp16 动态范围内可表示的数字。虽然人们可能希望缩放始终大于 1，但我们的 GradScaler 不会保证这一点，以保持性能。如果在使用 AMP/fp16 运行时遇到损失或梯度中的 NaN，请确保您的模型是兼容的。

2.1 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…) is equivalent to torch.amp.GradScaler(“cuda”, args…)

总结

API最后还列举了一些CUDA/CPU数据类型的可转换介绍，详细可访问链接