torch.utils.rename_privateuse1_backend

torch.utils.rename_privateuse1_backend(backend_name)

重新命名 privateuse1 後端裝置，使其更方便地在 PyTorch API 中用作裝置名稱。

步驟如下：

1. (在 C++) 中，為各種 torch 運算實作核心 (kernels)，並將它們註冊到 PrivateUse1 調度鍵 (dispatch key)。

2. (在 Python 中) 呼叫 torch.utils.rename_privateuse1_backend(“foo”)

現在您可以在 Python 中將 “foo” 作為一個普通的裝置字串使用。

注意：此 API 在每個程序中只能呼叫一次。嘗試在已設定後更改外部後端 (backend) 將導致錯誤。

注意(AMP)：如果您想在您的裝置上支援 AMP，您可以註冊一個自訂的後端模組。後端必須使用 torch._register_device_module("foo", BackendModule) 註冊一個自訂的後端模組。BackendModule 需要具有以下 API：

1. get_amp_supported_dtype() -> List[torch.dtype] 取得您的 “foo” 裝置在 AMP 中支援的資料類型 (dtypes)，或許 “foo” 裝置支援更多資料類型。

注意(random)：如果您想支援為您的裝置設定 seed，BackendModule 需要具有以下 API：

1. _is_in_bad_fork() -> bool 如果目前處於 bad_fork 狀態，則返回 True，否則返回 False。

2. manual_seed_all(seed int) -> None 為您的裝置設定用於產生隨機數字的 seed。

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.dev.org.tw/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")