torchmetrics: 一個用於分散式訓練與評估的可擴展 PyTorch 指標函式庫

它解決了什麼問題

TorchMetrics 為 PyTorch 應用程式提供了一種標準化的方式來計算機器學習指標。它消除了在多個批次（batch）與分散式裝置（例如多個 GPU 或節點）之間累積與同步指標時通常需要的樣板程式碼，確保結果具有可重現性與可擴展性。

運作方式

該函式庫提供兩種主要的指標計算方式：

• 基於模組的指標 (Module-based metrics)：這些指標的作用方式類似於 PyTorch 模組，透過維護內部狀態來自動追蹤並在不同批次之間累積數據。它們會自動處理多個裝置之間的同步，使其與 CPU、單 GPU 或多 GPU 設定相容。
• 功能性指標 (Functional metrics)：這些是簡單的 Python 函式，以張量（tensor）作為輸入並立即回傳指標值，而不維護狀態。

使用者也可以透過繼承 torchmetrics.Metric 並定義指標應如何更新其狀態與計算最終結果，來建立自定義指標。

對象是誰

它是為 PyTorch 開發者與機器學習工程師設計的，這些開發者需要追蹤不同領域（音訊、文本、圖像等）的模型效能，以及從事大規模分散式訓練的人員。

重點特色

• 豐富的函式庫：包含超過 100 種內建指標，涵蓋分類、回歸、分割、音訊、文本與多模態數據。
• 分散式支援：內建針對多裝置訓練的自動同步與累積功能。
• 可自定義：提供簡單的 API，可透過繼承基底類別來實作自定義指標。
• 視覺化：整合了繪圖支援，以視覺化指標隨時間變化的進度。