python - 使用 torch.autocast 时，如何强制各个图层 float 32

Question

我正在尝试以混合精度训练模型。但是，出于稳定性原因，我希望其中一些层完全精确。使用 torch.autocast 时如何强制单个层为 float32？特别是，我希望它是 onnx 可编译的。

是不是像:

with torch.cuda.amp.autocast(enabled=False, dtype=torch.float32):
    out = my_unstable_layer(inputs.float())

编辑:

看来这确实是官方的方法。查看torch docs .

Answer 1

我认为torch.autocast的动机是自动降低精度(而不是提高精度)。

如果你有functions that need a particular dtype ，你应该考虑使用，custom_fwd

import torch
@torch.cuda.amp.custom_fwd(cast_inputs=torch.complex128)
def get_custom(x):
    print('  Decorated function received', x.dtype)
def regular_func(x):
    print('  Regular function received', x.dtype)
    get_custom(x)

x = torch.tensor(0.0, dtype=torch.half, device='cuda')
with torch.cuda.amp.autocast(False):
    print('autocast disabled')
    regular_func(x)
with torch.cuda.amp.autocast(True):
    print('autocast enabled')
    regular_func(x)

autocast disabled
  Regular function received torch.float16
  Decorated function received torch.float16
autocast enabled
  Regular function received torch.float16
  Decorated function received torch.complex128

编辑:使用 torchscript

由于文档中的评论，我不确定您可以在多大程度上依赖它。然而，评论显然已经过时了。

这是一个示例，我在启用自动转换的情况下跟踪模型，将其卡住，然后使用它，并且值确实被转换为指定的类型

class Cast(torch.nn.Module):    
    @torch.cuda.amp.custom_fwd(cast_inputs=torch.float64)
    def forward(self, x):
        return x

with torch.cuda.amp.autocast(True):
    model = torch.jit.trace(Cast().eval(), x)
model = torch.jit.freeze(model)

x = torch.tensor(0.0, dtype=torch.half, device='cuda')
print(model(x).dtype)

torch.float64

但我建议您在将此方法用于正式应用之前对其进行验证。