torch.utils.generate_methods_for_privateuse1_backend

torch.utils.generate_methods_for_privateuse1_backend(for_tensor=True, for_module=True, for_packed_sequence=True, for_storage=False, unsupported_dtype=None)[源代码]

在将后端重命名为 privateuse1 之后,自动为自定义后端生成属性和方法。

默认情况下,与存储相关的方法不会自动生成。

当你为各种torch操作实现内核,并将其注册到PrivateUse1调度键时,可以调用函数torch.rename_privateuse1_backend("foo")来重命名你的后端名称。此时,你可以轻松地通过调用此函数来注册特定的方法和属性,例如torch.Tensor.foo()、torch.Tensor.is_foo、torch.Storage.foo()、torch.Storage.is_foo。

注意:建议使用通用函数(如检查设备是否相等或to(device=))。我们仅出于方便提供这些方法,它们会被“monkey patched”到对象上,因此不会正确进行类型化。对于Storage的generate方法,如果你需要支持稀疏数据存储,则需自行扩展其实现。

参数
  • for_tensor (bool) – 是否为 torch.Tensor 类注册相关的方法。

  • for_module (bool) – 是否为 torch.nn.Module 类注册相关的方法。

  • for_storage (bool) – 是否为 torch.Storage 类注册相关的方法。

  • unsupported_dtype (List[torch.dtype]) – 当需要生成存储方法时生效,表示该存储不支持 torch.dtype 类型。

示例:

>>> torch.utils.rename_privateuse1_backend("foo")
>>> torch.utils.generate_methods_for_privateuse1_backend()
# Then automatically generate backend-related attributes and methods.
>>> a = torch.tensor(2).foo()
>>> a.is_foo
>>> hasattr(torch.nn.Module, 'foo')
本页目录