speechbrain.utils.autocast 模块

该模块实现了与torch.autocast一起使用的实用工具和抽象，即自动混合精度。

Authors

  • Sylvain de Langen 2023

摘要

函数：

fwd_default_precision

用于前向方法的装饰器，默认情况下，禁用自动转换并将任何浮点张量参数转换为指定的数据类型（类似于torch.cuda.amp.custom_fwd）。

参考

speechbrain.utils.autocast.fwd_default_precision(fwd: Callable | None = None, cast_inputs: dtype | None = torch.float32)[source]

用于前向方法的装饰器，默认情况下，禁用自动转换并将任何浮点张量参数转换为指定的数据类型（类似于torch.cuda.amp.custom_fwd）。

wrapped forward 将获得一个额外的 force_allow_autocast 关键字参数。 当设置为 True 时，函数将忽略 cast_inputs 并且不会禁用自动转换，就像没有指定此装饰器一样。 （因此，模块可以指定默认的推荐精度，用户可以在需要时覆盖该行为。）

请注意，截至 PyTorch 2.1.1，这仅会影响CUDA AMP。 非 CUDA AMP 将不受影响，并且不会对输入张量进行类型转换！ 未来可能会通过此函数支持此用例。

当自动转换未激活时，此装饰器不会改变任何行为。

Parameters:

  • fwd (可选[可调用]) –

    要包装的函数。如果省略，则返回装饰器的部分应用，例如允许 new_decorator = fwd_default_precision(cast_inputs=torch.float32)。

    提醒：如果你直接装饰一个函数，这个参数已经隐式指定。

  • cast_inputs (可选[torch.dtype]) –

    如果不是 None（默认是 torch.float32），那么任何传递给包装函数的浮点输入将被转换为指定的类型。

    注意：当启用自动转换时，自动转换兼容操作的输出张量可能是自动转换数据类型。 禁用自动转换 而不 转换输入不会改变这一事实，因此即使在禁用自动转换的区域内也可能发生低精度操作，如果需要，此参数可以帮助避免这种情况。

Return type:

包装后的函数