作者:Maxim Lukiyanov - Microsoft 首席產品經理、Guoliang Hua - Microsoft 首席工程經理、Geeta Chauhan - Facebook 合作夥伴工程主管、Gisle Dankel - Facebook 技術主管

隨著 PyTorch 1.8.1 版本發布,我們很高興宣布 PyTorch Profiler – 全新改良的 PyTorch 效能除錯分析器。PyTorch Profiler 是由 Microsoft 和 Facebook 合作開發的開源工具,可針對大型深度學習模型進行精確且高效的效能分析和疑難排解。

分析和改善大型深度學習模型的效能是一項持續存在的挑戰,隨著模型尺寸增加,其重要性也日益提升。長期以來,由於缺乏可用的工具,PyTorch 使用者難以解決此挑戰。雖然有提供 GPU 硬體層級資訊的標準效能除錯工具,但卻遺漏了 PyTorch 特定操作的上下文。為了恢復遺失的資訊,使用者需要將多個工具組合在一起,或手動新增最少的關聯資訊才能理解資料。此外,還有 autograd profiler (torch.autograd.profiler),它可以擷取有關 PyTorch 操作的資訊,但無法擷取詳細的 GPU 硬體層級資訊,也無法提供視覺化支援。

全新的 PyTorch Profiler (torch.profiler) 是一個工具,它將這兩種類型的資訊結合在一起,然後建立體驗,充分發揮這些資訊的潛力。這個新的分析器收集 GPU 硬體和 PyTorch 相關資訊,將它們關聯起來,自動偵測模型中的瓶頸,並產生關於如何解決這些瓶頸的建議。分析器中的所有這些資訊都會在 TensorBoard 中視覺化呈現給使用者。新的 Profiler API 在 PyTorch 中原生支援,並提供迄今為止最簡單的體驗,使用者無需安裝任何額外套件即可分析模型,並透過新的 PyTorch Profiler 外掛程式立即在 TensorBoard 中查看結果。下方是 PyTorch Profiler 的螢幕截圖 - 自動瓶頸偵測。

開始使用

PyTorch Profiler 是 PyTorch autograd profiler 的下一個版本。它有一個新的模組命名空間 torch.profiler,但保持與 autograd profiler API 的相容性。Profiler 使用新的 GPU 分析引擎,該引擎使用 Nvidia CUPTI API 建構,並且能夠以高保真度擷取 GPU 核心事件。若要分析您的模型訓練迴圈,請將程式碼包裝在分析器上下文管理器中,如下所示。

 with torch.profiler.profile(
    schedule=torch.profiler.schedule(
        wait=2,
        warmup=2,
        active=6,
        repeat=1),
    on_trace_ready=tensorboard_trace_handler,
    with_stack=True
) as profiler:
    for step, data in enumerate(trainloader, 0):
        print("step:{}".format(step))
        inputs, labels = data[0].to(device=device), data[1].to(device=device)

        outputs = model(inputs)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        profiler.step()

schedule 參數可讓您限制包含在分析中的訓練步驟數量,以減少收集的資料量,並透過專注於重要內容來簡化視覺分析。tensorboard_trace_handler 會自動將分析結果儲存到磁碟,以便在 TensorBoard 中進行分析。

若要在 TensorBoard 中查看分析會話的結果,請安裝 PyTorch Profiler TensorBoard 外掛程式套件。

pip install torch_tb_profiler

Visual Studio Code 整合

Microsoft Visual Studio Code 是 Python 開發人員和資料科學家最受歡迎的程式碼編輯器之一。VS Code 的 Python 擴充功能最近將 TensorBoard 整合到程式碼編輯器中,包括對 PyTorch Profiler 的支援。安裝 VS Code 和 Python 擴充功能後,您可以透過鍵盤快捷鍵 CTRL + SHIFT + P (Mac 上為 CMD + SHIFT + P) 啟動命令面板,並輸入「Launch TensorBoard」命令,快速開啟 TensorBoard Profiler 外掛程式。

此整合具有內建的生命週期管理功能。如果您系統上沒有 TensorBoard 套件和 PyTorch Profiler 外掛程式套件 (將於四月中旬推出),VS Code 將自動安裝它們。VS Code 也會為您啟動 TensorBoard 程序,並自動尋找您目前目錄中的任何 TensorBoard 記錄檔。完成後,只需關閉標籤頁,VS Code 就會自動關閉程序。不再需要終端機視窗在您的系統上執行,以提供 TensorBoard UI 的後端!下方是在 TensorBoard 中執行的 PyTorch Profiler 追蹤檢視。

在此部落格中了解更多關於 VS Code 中 TensorBoard 支援的資訊。

意見回饋

請查看 PyTorch Profiler 文件,試用 Profiler 並告訴我們您的體驗。在 PyTorch 討論區提供您的意見回饋,或在 PyTorch GitHub 上提交問題。