torch.utils.rename_privateuse1_backend¶
- torch.utils.rename_privateuse1_backend(backend_name)[源代码]¶
重命名 privateuse1 后端设备,以便更方便地将其用作 PyTorch API 中的设备名称。
步骤如下:
(在 C++ 中) 为各种 torch 操作实现内核,并将它们注册到 PrivateUse1 分派键。
(在 Python 中) 调用 torch.utils.rename_privateuse1_backend(“foo”)
您现在可以在 Python 中将“foo”用作普通的设备字符串。
注意:此 API 每个进程只能调用一次。尝试在已设置外部后端后更改它会导致错误。
注意(AMP):如果您想在您的设备上支持 AMP,您可以注册一个自定义后端模块。后端必须使用
torch._register_device_module("foo", BackendModule)
注册一个自定义后端模块。BackendModule 需要具有以下 API:get_amp_supported_dtype() -> List[torch.dtype]
获取您“foo”设备在 AMP 中支持的数据类型,也许“foo”设备支持一种或多种数据类型。
注意(随机数):如果您想支持为您的设备设置种子,BackendModule 需要具有以下 API:
_is_in_bad_fork() -> bool
如果现在处于 bad_fork 状态,则返回True
,否则返回False
。manual_seed_all(seed int) -> None
设置为您的设备生成随机数的种子。device_count() -> int
返回可用的“foo”数量。get_rng_state(device: Union[int, str, torch.device] = 'foo') -> Tensor
返回一个 ByteTensor 列表,表示所有设备的随机数状态。set_rng_state(new_state: Tensor, device: Union[int, str, torch.device] = 'foo') -> None
设置指定“foo”设备的随机数生成器状态。
还有一些常用函数:
is_available() -> bool
返回一个布尔值,指示“foo”当前是否可用。current_device() -> int
返回当前选定设备的索引。
有关更多详细信息,请参阅 https://pytorch.ac.cn/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")