Механизм автозагрузки

Предыстория

Предложение Autoload Device Extension в PyTorch направлено на улучшение поддержки различных аппаратных устройств-бэкендов, особенно тех, которые реализованы как внешние расширения (не входящие в основную кодовую базу PyTorch). В настоящее время пользователям приходится вручную импортировать или загружать такие расширения для конкретных устройств, что усложняет работу и увеличивает когнитивную нагрузку.

В отличие от этого, устройства, встроенные в PyTorch (официально поддерживаемые устройства), интегрированы без проблем — пользователям не нужны дополнительные импорты или действия. Цель автозагрузки — сделать внешние устройства такими же простыми в использовании, чтобы пользователи могли следовать стандартной модели программирования устройств PyTorch без явной загрузки или изменений в коде. Благодаря этому существующие приложения PyTorch смогут работать на новых устройствах без каких-либо изменений, что сделает поддержку оборудования удобнее для пользователей и снизит барьеры для внедрения.

Дополнительные сведения о предпосылках автозагрузки см. в соответствующем RFC.

Проектирование

Основная идея автозагрузки — использовать механизм обнаружения плагинов Python (точки входа), чтобы PyTorch автоматически загружал внешние расширения устройств при импорте torch, без необходимости явного импорта со стороны пользователя.

Дополнительные инструкции по проектированию автозагрузки см. в разделе Как это работает.

Реализация

В этом руководстве в качестве нового внешнего устройства будет использоваться OpenReg. Вы узнаете, как включить и использовать механизм автозагрузки.

 1    setup(
 2        packages=find_packages(),
 3        package_data=package_data,
 4        ext_modules=ext_modules,
 5        cmdclass={
 6            "clean": BuildClean,  # type: ignore[misc]
 7        },
 8        include_package_data=False,
 9        entry_points={
10            "torch.backends": [
11                "torch_openreg = torch_openreg:_autoload",
12            ],
13        },
14    )

1def _autoload():
2    # It is a placeholder function here to be registered as an entry point.
3    pass
4
5

Результат

После настройки точки входа и бэкенда соберите и установите свой бэкенд. Теперь новый ускоритель можно использовать без его явного импорта.

>>> import torch
>>> import torch_openreg
>>> torch.tensor(1, device="openreg")
tensor(1, device='openreg:0')

>>> import torch # Automatically import torch_openreg
>>>
>>> torch.tensor(1, device="openreg")
tensor(1, device='openreg:0')