torch.utils.backend_registration.rename_privateuse1_backend

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

Переименовывает устройство бэкенда privateuse1, чтобы его было удобнее использовать в качестве имени устройства в API PyTorch.

Шаги:

1. (На C++) реализуйте ядра для различных операций torch и зарегистрируйте их для ключа диспетчеризации PrivateUse1.
2. (На Python) вызовите torch.utils.rename_privateuse1_backend(“foo”)

Теперь вы можете использовать “foo” как обычную строку устройства в Python.

Примечание: этот API можно вызвать только один раз в рамках процесса. Попытка изменить внешний бэкенд после его установки приведёт к ошибке.

Примечание (AMP): если вы хотите обеспечить поддержку AMP на своём устройстве, можно зарегистрировать пользовательский модуль бэкенда. Бэкенд должен зарегистрировать пользовательский модуль бэкенда с помощью torch._register_device_module("foo", BackendModule). BackendModule должен предоставлять следующие API:

1. get_amp_supported_dtype() -> List[torch.dtype] получить поддерживаемые типы данных для устройства “foo” в AMP; возможно, устройство “foo” поддерживает ещё один тип данных.

Примечание (случайные числа): если вы хотите иметь возможность задавать начальное значение генератора для своего устройства, BackendModule должен предоставлять следующие API:

1. _is_in_bad_fork() -> bool Возвращает True, если сейчас действует bad_fork, иначе возвращает 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.org/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")