自動載入樹外擴充功能¶

建立於：2024 年 10 月 10 日 | 最後更新：2024 年 10 月 10 日 | 最後驗證：2024 年 10 月 10 日

擴充功能自動載入機制使 PyTorch 能夠自動載入樹外後端擴充功能，而無需明確的 import 語句。這個功能對使用者來說是有益的，因為它可以增強他們的使用體驗，並且讓他們能夠遵循熟悉的 PyTorch 裝置程式設計模型，而無需明確載入或匯入裝置特定的擴充功能。此外，它還能夠讓現有的 PyTorch 應用程式輕鬆採用樹外裝置，而無需修改任何程式碼。更多詳細資訊，請參閱 [RFC] Autoload Device Extension。

您將學到什麼

如何在 PyTorch 中使用樹外擴充功能自動載入
檢閱 Intel Gaudi HPU、Huawei Ascend NPU 的範例

先決條件

PyTorch v2.5 或更高版本

注意

此功能預設為啟用，可以使用 export TORCH_DEVICE_BACKEND_AUTOLOAD=0 停用。如果您收到類似這樣的錯誤：「Failed to load the backend extension」，這個錯誤與 PyTorch 無關，您應該停用此功能並尋求樹外擴充功能維護者的協助。

如何將此機制應用於樹外擴充功能？¶

例如，假設您有一個名為 foo 的後端，以及一個名為 torch_foo 的對應套件。請確保您的套件與 PyTorch 2.5 或更高版本相容，並且在其 __init__.py 檔案中包含以下程式碼片段

def _autoload():
    print("Check things are working with `torch.foo.is_available()`.")

然後，您唯一需要做的就是在您的 Python 套件中定義一個進入點

setup(
    name="torch_foo",
    version="1.0",
    entry_points={
        "torch.backends": [
            "torch_foo = torch_foo:_autoload",
        ],
    }
)

現在，您只需新增 import torch 語句即可匯入 torch_foo 模組，而無需新增 import torch_foo

>>> import torch
Check things are working with `torch.foo.is_available()`.
>>> torch.foo.is_available()
True

在某些情況下，您可能會遇到循環匯入的問題。以下範例示範了如何解決這些問題。

範例¶

在本範例中，我們將使用 Intel Gaudi HPU 和 Huawei Ascend NPU 來確定如何使用自動載入功能將您的樹外擴充功能與 PyTorch 整合。

habana_frameworks.torch 是一個 Python 套件，讓使用者能夠透過使用 PyTorch HPU 裝置金鑰在 Intel Gaudi 上執行 PyTorch 程式。

habana_frameworks.torch 是 habana_frameworks 的一個子模組，我們在 habana_frameworks/setup.py 中新增一個進入點到 __autoload()

setup(
    name="habana_frameworks",
    version="2.5",
+   entry_points={
+       'torch.backends': [
+           "device_backend = habana_frameworks:__autoload",
+       ],
+   }
)

在 habana_frameworks/init.py 中，我們使用一個全域變數來追蹤我們的模組是否已載入

import os

is_loaded = False  # A member variable of habana_frameworks module to track if our module has been imported

def __autoload():
    # This is an entrypoint for pytorch autoload mechanism
    # If the following condition is true, that means our backend has already been loaded, either explicitly
    # or by the autoload mechanism and importing it again should be skipped to avoid circular imports
    global is_loaded
    if is_loaded:
        return
    import habana_frameworks.torch

在 habana_frameworks/torch/init.py 中，我們透過更新全域變數的狀態來防止循環匯入

import os

# This is to prevent torch autoload mechanism from causing circular imports
import habana_frameworks

habana_frameworks.is_loaded = True

torch_npu 讓使用者能夠在 Huawei Ascend NPU 上執行 PyTorch 程式，它利用 PrivateUse1 裝置金鑰，並將裝置名稱公開為 npu 給終端使用者。

我們在 torch_npu/setup.py 中定義一個進入點

setup(
    name="torch_npu",
    version="2.5",
+   entry_points={
+       'torch.backends': [
+           'torch_npu = torch_npu:_autoload',
+       ],
+   }
)

與 habana_frameworks 不同，torch_npu 使用環境變數 TORCH_DEVICE_BACKEND_AUTOLOAD 來控制自動載入程序。例如，我們將其設定為 0 以停用自動載入，以防止循環匯入

# Disable autoloading before running 'import torch'
os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'

import torch

運作方式¶

自動載入是基於 Python 的 Entrypoints 機制實現的。我們會探索並載入 torch/__init__.py 中所有由樹外擴充功能定義的特定進入點。

如上所示，在安裝 torch_foo 之後，當載入您定義的進入點時，您可以匯入您的 Python 模組，然後您可以在呼叫它時執行一些必要的工作。

請參閱此 Pull Request 中的實作：[RFC] Add support for device extension autoloading。

結論¶

在本教學中，我們了解了 PyTorch 中的樹外擴充功能自動載入機制，該機制會自動載入後端擴充功能，從而消除了新增額外 import 語句的需要。我們還了解了如何透過定義進入點將此機制應用於樹外擴充功能，以及如何防止循環匯入。我們還檢閱了一個關於如何將自動載入機制與 Intel Gaudi HPU 和 Huawei Ascend NPU 搭配使用的範例。

自動載入樹外擴充功能¶

如何將此機制應用於樹外擴充功能？¶

範例¶

運作方式¶

結論¶

文件

教學

資源