Python PyTorch autocast用法及代码示例

本文简要介绍python语言中 torch.autocast 的用法。

用法:

class torch.autocast(device_type, enabled=True, **kwargs)

参数：

• device_type(string,必需的) -是否使用‘cuda’或‘cpu’设备

• enabled(bool,可选的,默认=真) -是否应在区域中启用自动投射。

• dtype(torch_dpython:类型,可选的) -是使用 torch.float16 还是 torch.bfloat16。

• cache_enabled(bool,可选的,默认=真) -是否应该启用 autocast 中的权重缓存。

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

在这些区域中，ops 在 autocast 选择的 op-specific dtype 中运行，以提高性能同时保持准确性。有关详细信息，请参阅 Autocast Op 参考。

进入autocast-enabled 区域时，张量可以是任何类型。使用自动投射时，不应在模型或输入上调用 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 autocast():
        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):
    ...
    @autocast()
    def forward(self, input):
        ...

在 autocast-enabled 区域中产生的浮点张量可能是 float16 。返回到 autocast-disabled 区域后，将它们与不同 dtype 的浮点张量一起使用可能会导致类型不匹配错误。如果是这样，请将自动投射区域中生成的张量转换回 float32(或其他 dtype，如果需要)。如果来自自动投射区域的张量已经是 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 autocast():
    # 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())

中央处理器示例：

# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cpu")
b_float32 = torch.rand((8, 8), device="cpu")
c_float32 = torch.rand((8, 8), device="cpu")
d_float32 = torch.rand((8, 8), device="cpu")

with autocast(dtype=torch.bfloat16, device_type="cpu"):
    # torch.mm is on autocast's list of ops that should run in bfloat16.
    # Inputs are float32, but the op runs in bfloat16 and produces bfloat16 output.
    # No manual casts are required.
    e_bfloat16 = torch.mm(a_float32, b_float32)
    # Also handles mixed input types
    f_bfloat16 = torch.mm(d_float32, e_bfloat16)

# After exiting autocast, calls f_float16.float() to use with d_float32
g_float32 = torch.mm(d_float32, f_bfloat16.float())

autocast-enabled 区域中的类型不匹配错误是一个错误；如果这是您观察到的，请提出问题。

autocast(enabled=False) 子区域可以嵌套在 autocast-enabled 区域中。例如，如果您想强制子区域在特定的 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 autocast():
    e_float16 = torch.mm(a_float32, b_float32)
    with autocast(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)。