torch.utils.rename_privateuse1_backend

torch.utils.rename_privateuse1_backend(backend_name)

将 privateuse1 后端设备重命名，以便在 PyTorch API 中更方便地用作设备名称。

步骤如下：

1. （在 C++ 中）为各种 torch 操作实现内核，并将它们注册到 PrivateUse1 dispatch key。

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] 获取 AMP 中您的“foo”设备支持的 dtype，也许“foo”设备支持更多 dtype。

注意（random）：如果您想为设备设置种子，BackendModule 需要具有以下 API：

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

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

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

4. get_rng_state(device: Union[int, str, torch.device] = 'foo') -> Tensor 返回一个表示所有设备随机数状态的 ByteTensor 列表。

5. set_rng_state(new_state: Tensor, device: Union[int, str, torch.device] = 'foo') -> None 设置指定“foo”设备的随机数生成器状态。

还有一些通用函数：

1. is_available() -> bool 返回一个布尔值，指示“foo”当前是否可用。

2. current_device() -> int 返回当前选定设备的索引。

有关更多详细信息，请参阅 https://pytorch.ac.cn/tutorials/advanced/extend_dispatcher.html#get-a-dispatch-key-for-your-backend。现有示例请参阅 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")