部落格

介紹 PyTorch Profiler – 新改進的效能工具

隨著 PyTorch 1.8.1 版本的釋出，我們很高興地推出 PyTorch Profiler —— 專為 PyTorch 設計的全新且改進的效能除錯分析器。PyTorch Profiler 是微軟與 Facebook 合作開發的開源工具，旨在為大規模深度學習模型提供準確、高效的效能分析和故障排查功能。

分析和提升大規模深度學習模型的效能是一個持續存在的挑戰，且隨著模型規模的增長，其重要性日益凸顯。長期以來，由於缺乏合適的工具，PyTorch 使用者在解決這一挑戰時感到非常困難。現有的標準效能除錯工具雖然能提供 GPU 硬體層面的資訊，卻缺失了 PyTorch 運算元層面的上下文。為了獲取這些缺失的資訊，使用者通常需要組合多種工具，或者手動新增關聯資訊來解讀資料。此外，雖然原有的自動求導分析器 (torch.autograd.profiler) 可以捕獲 PyTorch 運算元的相關資訊，但它無法獲取詳細的 GPU 硬體級資訊，也無法提供視覺化支援。

全新的 PyTorch Profiler (torch.profiler) 將上述兩類資訊整合在一起，並打造了能夠充分發揮這些資料價值的體驗。該分析器不僅能收集 GPU 硬體和 PyTorch 相關資訊，還能將兩者關聯起來，自動識別模型中的效能瓶頸，並生成解決這些瓶頸的建議。分析器收集的所有資訊都會透過 TensorBoard 為使用者呈現。全新的 Profiler API 在 PyTorch 中是原生支援的，提供了迄今為止最簡便的使用體驗：使用者無需安裝任何額外包即可分析模型，並透過全新的 PyTorch Profiler 外掛直接在 TensorBoard 中檢視結果。下方為 PyTorch Profiler 自動瓶頸檢測功能的截圖。

入門指南

PyTorch Profiler 是 PyTorch 自動求導分析器的下一代版本。它採用了全新的模組名稱空間 torch.profiler，但保持了與自動求導分析器 API 的相容性。該分析器使用了基於 Nvidia CUPTI API 構建的全新 GPU 分析引擎，能夠以高保真度捕獲 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 外掛包（4 月中旬釋出），VS Code 將自動為你安裝。VS Code 還會自動為你啟動 TensorBoard 程序，並自動掃描當前目錄下的任何 TensorBoard 日誌檔案。當你完成工作時，只需關閉標籤頁，VS Code 就會自動結束該程序。無需再讓終端視窗在後臺執行來提供 TensorBoard UI 的後端服務了！下方是正在 TensorBoard 中執行的 PyTorch Profiler 追蹤檢視 (Trace View)。

欲瞭解更多關於 VS Code 中 TensorBoard 支援的資訊，請參閱此部落格。

反饋

請檢視 PyTorch Profiler 文件，嘗試使用該分析器並將你的使用體驗反饋給我們。歡迎在 PyTorch 論壇提交反饋，或在 PyTorch GitHub 上提交問題。