PyTorch 大規(guī)模部署的功能

原文： https://pytorch.org/docs/stable/notes/large_scale_deployments.html

本說明討論了一些擴(kuò)展點和技巧，這些擴(kuò)展點和技巧在較大的系統(tǒng)中運行 PyTorch 或在較大的組織中使用 PyTorch 操作多個系統(tǒng)時可能有用。

它不涉及將模型部署到生產(chǎn)中的主題。 檢查 torch.jit 或相應(yīng)的教程之一。

該注釋假定您是從組織中的源代碼構(gòu)建 PyTorch 的，或者能夠靜態(tài)鏈接使用 PyTorch 時要加載的其他代碼。 因此，許多鉤子都公開為 C ++ API，可以在集中式位置(例如， 在靜態(tài)初始化代碼中。

整個機(jī)隊的運營商配置文件

PyTorch 帶有torch.autograd.profiler，能夠測量各個操作員根據(jù)需要花費的時間。 可以使用相同的機(jī)制對運行 PyTorch 的任何進(jìn)程進(jìn)行“始終在線”測量。 收集有關(guān)在給定進(jìn)程中或在整個機(jī)器組中運行的 PyTorch 工作負(fù)載的信息可能很有用。

可以使用torch::autograd::profiler::pushCallback添加用于任何操作員調(diào)用的新回調(diào)。 掛鉤將使用描述調(diào)用上下文的torch::autograd::profiler::RecordFunction結(jié)構(gòu)調(diào)用(例如<cite>名稱</cite>）。 如果啟用，RecordFunction::inputs()包含表示為torch::IValue變量類型的函數(shù)的參數(shù)。 注意，輸入日志記錄相對昂貴，因此必須顯式啟用。

操作員回調(diào)還可以訪問at::getThreadLocalDebugInfo()接口，該接口返回指向包含調(diào)試信息的結(jié)構(gòu)的指針。 該調(diào)試信息應(yīng)該與相應(yīng)的at::setThreadLocalDebugInfo(debug_info)調(diào)用一起設(shè)置。 調(diào)試信息通過前向傳播(包括異步fork任務(wù)）和后向傳播進(jìn)行傳播，對于將有關(guān)執(zhí)行環(huán)境的一些額外信息(例如，模型 ID）從應(yīng)用程序的高層傳遞到操作員回調(diào)非常有用。

調(diào)用回調(diào)會增加一些開銷，因此通常隨機(jī)抽樣操作員調(diào)用很有用。 可以在每個回調(diào)的基礎(chǔ)上啟用 <cite>torch :: autograd :: profiler :: setSamplingProbability</cite> 指定的全局采樣率。

請注意，pushCallback和setSamplingProbability不是線程安全的，只有在沒有運行 PyTorch 運算符時才能調(diào)用。 通常，在初始化過程中一次調(diào)用它們是一個好主意。

這是一個例子：

// Called somewhere in the program beginning
void init() {
    // Sample one in a hundred operator runs randomly
    torch::autograd::setSamplingProbability(0.01);
    pushCallback(
        &onFunctionEnter,
        &onFunctionExit,
        /* needs_inputs */ true,
        /* sampled */ true
    );
}


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 使用記錄

在更廣泛的生態(tài)系統(tǒng)中運行時(例如在托管的工作計劃程序中），跟蹤哪些二進(jìn)制文件調(diào)用特定的 PyTorch API 通常很有用。 在幾個重要的 API 點注入了簡單的工具，這些工具會觸發(fā)給定的回調(diào)。 由于通常在一次性 python 腳本中調(diào)用 PyTorch，因此對于每個 API 的給定進(jìn)程，回調(diào)僅觸發(fā)一次。

c10::SetAPIUsageHandler可用于注冊 API 使用情況檢測處理程序。 傳遞的參數(shù)將是“ api key”，用于標(biāo)識使用的點，例如，用于 PyTorch 擴(kuò)展名導(dǎo)入的python.import或觸發(fā) TorchScript 編譯的torch.script.compile。

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

開發(fā)人員注意：可以在代碼中使用 C ++中的C10_LOG_API_USAGE_ONCE("my_api")或 Python 中的torch._C._log_api_usage_once("my.api")添加新的 API 觸發(fā)點。

將元數(shù)據(jù)附加到已保存的 TorchScript 模型中

TorchScript 模塊可以保存為存檔文件，該文件將序列化的參數(shù)和模塊代碼捆綁為 TorchScript(請參見 torch.jit.save())。 將附加信息與模型捆綁在一起通常很方便，例如，模型生產(chǎn)者或輔助工件的描述。

可以通過將_extra_files<>參數(shù)傳遞給 torch.jit.save() 和torch::jit::load<>在存儲過程中存儲和檢索任意二進(jìn)制 Blob 來實現(xiàn)。 由于 TorchScript 文件是常規(guī)的 ZIP 存檔，因此額外的信息將作為常規(guī)文件存儲在存檔的extra/<>目錄中。

還有一個全局掛鉤，可將其他文件附加到當(dāng)前流程中生成的任何 TorchScript 存檔中。 用生產(chǎn)者元數(shù)據(jù)標(biāo)記模型可能很有用，類似于由數(shù)碼相機(jī)產(chǎn)生的 JPEG 元數(shù)據(jù)。 用法示例如下所示：

SetExportModuleExtraFilesHook([](const script::Module&) {
    script::ExtraFilesMap files;
    files["producer_info.json"] = "{\"user\": \"" + getenv("USER") + "\"}";
    return files;
});

構(gòu)建環(huán)境注意事項

TorchScript 的編譯需要使用 python 的inspect.getsource<>調(diào)用，因此必須有權(quán)訪問原始 python 文件。 在某些生產(chǎn)環(huán)境中，可能需要顯式部署.py文件以及預(yù)編譯的.pyc文件。

通用擴(kuò)展點

PyTorch API 通常是松散耦合的，很容易用專用版本替換組件。 常見的擴(kuò)展點包括：

• 使用 C ++實現(xiàn)的自定義運算符-有關(guān)更多詳細(xì)信息，請參見教程。
• 自定義數(shù)據(jù)讀取通?？梢酝ㄟ^調(diào)用相應(yīng)的 python 庫直接集成。 torch.utils.data 的現(xiàn)有功能可以通過擴(kuò)展 Dataset 或  IterableDataset 加以利用。