大規模部署功能

創建於: 2019年7月24日 | 最後更新於: 2025年7月15日

本文件討論了在 PyTorch 運行於更大的系統內或在大型組織中使用 PyTorch 執行多個系統時可能有用的一些擴充套件點和技巧。

本文件假設您要麼從原始碼構建 PyTorch，要麼能夠靜態連結額外的程式碼，以便在 PyTorch 使用時載入。因此，許多鉤子都作為 C++ API 公開，可以在中心化位置（例如，靜態初始化程式碼）一次性觸發。

全叢集運算元分析

PyTorch 提供了 torch.autograd.profiler，能夠按需測量單個運算元所花費的時間。可以使用相同的機制對執行 PyTorch 的任何程序進行“始終開啟”的測量。這可能有助於收集給定程序或整個機器集上 PyTorch 工作負載的資訊。

可以使用 torch::addGlobalCallback 為任何運算元呼叫新增新的回撥。鉤子將使用 torch::RecordFunction 結構體進行呼叫，該結構體描述了呼叫上下文（例如，name）。如果啟用了輸入日誌記錄，RecordFunction::inputs() 將包含表示為 torch::IValue 變體型別的函式引數。請注意，輸入日誌記錄成本相對較高，因此必須顯式啟用。

運算元回撥還可以訪問 c10::ThreadLocalDebugInfo::get() 介面，該介面返回一個指向儲存除錯資訊的結構體的指標。此除錯資訊可以透過使用 at::DebugInfoGuard 物件提前設定。除錯資訊會透過前向（包括非同步 fork 任務）和後向傳播，並可用於將有關執行環境的額外資訊（例如，模型 ID）從應用程式的高層傳遞到底層運算元回撥。

呼叫回撥會增加一些開銷，因此通常最好僅隨機抽樣運算元呼叫。這可以透過在 torch::addGlobalCallback 中傳遞可選的取樣率來按回調進行啟用。

請注意，addGlobalCallback 不是執行緒安全的，只能在沒有 PyTorch 運算元執行時呼叫。通常，在初始化期間呼叫它們一次是個好主意。

示例如下：

// Called somewhere in the program beginning
void init() {
    // Sample one in a hundred operator runs randomly
    addGlobalCallback(
      RecordFunctionCallback(
        &onFunctionEnter,
        &onFunctionExit)
      .needsInputs(true)
      .samplingProb(0.01)
    );
    // Note, to enable observers in the model calling thread,
    // call enableRecordFunction() in the thread before running a model
}

void onFunctionEnter(const RecordFunction& fn) {
    std::cerr << "Before function " << fn.name()
              << " with " << fn.inputs().size() << " inputs" << std::endl;
}

void onFunctionExit(const RecordFunction& fn) {
    std::cerr << "After function " << fn.name();
}

API 使用日誌記錄

在更廣泛的生態系統中執行時，例如在託管作業排程器中，跟蹤哪些二進位制檔案呼叫了特定的 PyTorch API 通常很有用。在幾個重要的 API 點注入了一個簡單的插樁，用於觸發給定的回撥。由於 PyTorch 通常在一次性 Python 指令碼中呼叫，因此對於每個 API，回撥在給定程序中最多隻觸發一次。

可以使用 c10::SetAPIUsageHandler 註冊 API 使用插樁處理程式。傳遞的引數將是一個“api key”，用於標識使用的點，例如，用於 PyTorch 擴充套件匯入的 python.import。

SetAPIUsageLogger([](const std::string& event_name) {
    std::cerr << "API was used: " << event_name << std::endl;
});

給開發者的說明：可以使用 C++ 中的 C10_LOG_API_USAGE_ONCE("my_api") 或 Python 中的 torch._C._log_api_usage_once("my.api") 在程式碼中新增新的 API 觸發點。

通用擴充套件點

PyTorch API 通常是鬆散耦合的，很容易用專用版本替換某個元件。通用擴充套件點包括：

• 用 C++ 實現的自定義運算元 - 有關更多詳細資訊，請參閱教程。

• 自定義資料讀取通常可以透過呼叫相應的 Python 庫直接整合。透過擴充套件 Dataset 或 IterableDataset，可以利用 torch.utils.data 的現有功能。