torch.utils.rename_privateuse1_backend

torch.utils.rename_privateuse1_backend(backend_name)[source][source]

将私有用途 1 后端设备重命名，使其在 PyTorch API 中使用时更加方便。

步骤如下：

1. （在 C++中）实现各种 torch 操作的内核，并将它们注册到 PrivateUse1 调度键。

2. （在 python 中）调用 torch.utils.rename_privateuse1_backend(“foo”)

现在您可以在 Python 中将“foo”用作普通设备字符串。

注意：此 API 每个进程只能调用一次。在已设置之后尝试更改外部后端将导致错误。

注意（AMP）：如果您想在您的设备上支持 AMP，您可以注册一个自定义后端模块。后端必须使用 torch._register_device_module("foo", BackendModule) 注册自定义后端模块。BackendModule 需要以下 API：

1. get_amp_supported_dtype() -> List[torch.dtype] 获取您的“foo”设备在 AMP 上支持的 dtype，也许“foo”设备还支持一个额外的 dtype。

Note(随机): 如果您想支持为您的设备设置种子，后端模块需要以下 API：

1. _is_in_bad_fork() -> bool 如果当前处于 bad_fork，则返回 True ，否则返回 False 。

2. manual_seed_all(seed int) -> None 设置用于生成随机数的种子。

3. device_count() -> int 返回可用的“foo”数量。

4. 返回所有设备上随机数状态的 ByteTensor 列表。

5. 设置指定“foo”设备的随机数生成器状态。

并且有一些常用函数：

1. 返回一个 bool 值，指示“foo”当前是否可用。

2. 返回当前选中设备的索引。

更多详情请参阅 https://pytorch.org/tutorials/advanced/extend_dispatcher.html#get-a-dispatch-key-for-your-backend。有关现有示例，请参阅 https://github.com/bdhirsh/pytorch_open_registration_example。

示例：

>>> torch.utils.rename_privateuse1_backend("foo")
# This will work, assuming that you've implemented the right C++ kernels
# to implement torch.ones.
>>> a = torch.ones(2, device="foo")