多線程優化數據加載效率

此問題的傳統解決方案之一是擴展處理并并行化用戶編寫的函數。在這種方法中，用戶創建自定義算法，而系統則負責在同時計算任務的多個工作進程中擴展其執行。這就是 torch.DataLoader 發揮作用的地方。

本文記錄了我們通過從進程切換到線程來優化 torch.DataLoader 的實驗。這項探索之所以成為可能，是因為 Python 不斷努力刪除 GIL，使我們能夠重新思考深度學習工作流程中的并行性，并探索新的性能優化。

什么是 torch.DataLoader？工作原理是什么？

torch.DataLoader 是 PyTorch 中的基礎工具，有助于在深度學習應用中加載數據。它在管理數據輸入模型的方式方面發揮著關鍵作用，可確保流程高效且有效。

torch.DataLoader 的重要特性是，它能夠并行化加載過程，這在處理大型數據集時至關重要。

這種并行化通常通過創建多個工作進程來實現，每個進程負責加載部分數據，這些進程并行運行，從而能夠在訓練模型的同時加載和預處理數據。

并行性對于保持穩定的GPU數據流、盡量減少空閑時間和盡量提高資源利用率尤為重要。

可怕的 GIL

torch.DataLoader 使用進程來并行化數據加載任務，這種方法直接源于 Python 架構的一個基本方面，即全局解釋器鎖（GIL）。

GIL 是一種互斥體，可防止多個原生線程在 CPython (最廣泛使用的 Python 實現) 中同時執行 Python 字節碼。這鎖的引入目的是簡化內存管理，并通過在多個線程試圖同時訪問或修改 Python 對象時防止出現競爭條件，以確保線程安全。

雖然 GIL 使 Python 的內存管理變得簡單，并有助于避免復雜的并發錯誤，但它也施加了一個重大限制：Python 線程并非真正的并行。

在受 CPU 限制的任務中，處理能力是瓶頸，線程不得不輪流運行，導致性能不佳。這就是為什么 torch.DataLoader 使用進程而不是線程的原因。每個進程都在自己的內存空間中運行，完全繞過 GIL，并允許在多核處理器上真正并行執行。

當然，GIL 的影響并非完全是負面的。它通過減少開發者對線程安全的關注來簡化 Python 程序的開發，這也是 Python 如此受歡迎的原因之一。

另一方面，GIL 可能會成為 CPU 受限和多線程應用程序的瓶頸，因為它阻礙了多核系統的充分利用。這種權衡在 Python 社區中引發了關于其優缺點的持續爭論。

線程交換進程

隨著近期的發展，GIL 將在即將推出的 Python 版本中刪除。這為 Python 應用程序（包括深度學習）中的并行性開辟了新的可能性。

我們的一個關鍵想法是嘗試將 torch.DataLoader 中基于進程的并行與基于線程的并行交換（圖 1）。

1. ?

使用線程代替進程有幾個潛在優勢。線程通常比進程輕，從而加快上下文切換速度，并降低內存開銷。

然而，線程也帶來了它自己的挑戰，特別是在確保線程安全和避免死鎖等問題方面。

我們實施了基于線程的 torch.DataLoader 版本來探索這些可能性。結果很有趣，并且表明，在某些情況下，線程可以成為進程的可行替代方案。

基于線程的數據加載結果

為了評估在 torch.DataLoader 中使用線程替換進程對性能的影響，我們在不同的數據處理場景中進行了一系列實驗。結果突出了基于線程的并行性的潛力和局限性。

使用 nvImageCodec 進行圖像解碼

在使用 nvImageCodec 的圖像解碼場景中，出現了使用線程的最引人注目的案例之一。在這種場景中，與傳統的基于進程的方法相比，使用線程可顯著提高速度。

基準測試詳細信息：EPYC 9654 | H100 | 批量大小：512 | 圖像大小：640 x 408（JPEG）

實現這一改進的主要原因是 CUDA 上下文切換的減少。在切換上下文時，進程會引入更高的開銷，這會導致嚴重的延遲，尤其是在 GPU 加速的工作負載中。

另一方面，線程可以減少這種開銷，從而實現更快、更高效的執行。

使用 Pillow 進行圖像解碼

與 nvImageCodec 不同，我們對 Pillow (一種廣泛使用的 Python 圖像庫) 的實驗表明，線程方法比基于進程的方法略慢。

基準測試詳細信息：EPYC 9654 | 批量大小：512 | 圖像大小：640 x 408（JPEG）

這里的關鍵區別在于全局狀態的管理方式。Pillow 的操作涉及對存儲在字典中的全局狀態數據的頻繁訪問。當多個線程同時訪問這些共享數據時，當前的實現依賴于原子來安全地管理這些操作。

然而，atomics 可能會成為爭用的瓶頸，與每個工作者都有自己的隔離狀態的單獨進程相比，這會導致性能降低。

由于這一瓶頸，我們在 discuss.python.org 上發起了一場討論，重新探討凍結數據類型的想法，這可以通過實現更高效的讀取訪問而無需昂貴的原子來幫助緩解這些性能問題。

綜合結果：nvImageCodec 與 Pillow

為了更好地顯示性能差異，我們將 nvImageCodec 和 Pillow 場景的結果合并到一張圖表中（圖 4）。

基準測試詳細信息：EPYC 9654 | H100 | 批量大小：512 | 圖像大小：640 x 408（JPEG）

這種比較清楚地表明了兩種方法之間的鮮明對比：

• nvImageCodec：線程的性能明顯優于進程，這表明在依賴 CUDA 的 GPU 密集型任務中，線程方法非常有利。
• Pillow：進程仍然保持著微小的優勢，這表明涉及共享狀態的任務可能無法從線程中獲益。

這些發現強調，在基于 GPU 的場景中，移除 GIL 可以立即顯著提高速度。然而，隨著 Python 邁出了進入自由線程領域的第一步，我們應該更加努力地引入新的工具和概念，充分利用硬件功能并充分發揮語言的潛力。

基于線程的 Torch.DataLoader 的優缺點

雖然基于線程的 torch.DataLoader 在某些情況下表現出明顯優勢，但務必要權衡利弊。

優勢顯而易見：

• 開銷更低：線程的資源密集程度低于進程，因此內存占用率更低，上下文切換速度更快。
• 在某些情況下提供更好的性能：正如 nvImageCodec 實驗所示，線程可以減少同步開銷，從而提高整體性能。

缺點如下：

• 線程安全：確保代碼的線程安全可能具有挑戰性，尤其是在復雜的數據管道中。對于線程，總是存在更高的死鎖風險，這可能會中斷整個數據加載過程。
• 廣泛的同步：通常情況下，線程必須比進程更頻繁地同步。實現基于線程的執行需要在開發過程中進行更多的審查。
• 遷移現有實現：自由線程的 Python 生態系統尚處于開發的早期階段。調整深度學習項目具有的大量依賴項需要一些時間。

結束語

刪除 GIL 為優化 Python 中的深度學習工作流程帶來了新的機會。我們對基于線程的 torch.DataLoader 的探索表明，每當工作者實現涉及 GPU 處理時，它都是一個有益的方法。

然而，對于 CPU 操作而言，由于對數據結構的并行讀取訪問效率低下，性能往往會出現瓶頸，我們希望在未來能夠解決這一問題。

隨著 Python 的不斷發展，深度學習中數據加載的格局必將發生變化，我們很高興能走在這些發展的前沿。

如果您有興趣詳細了解我們使用自由線程 Python 進行的實驗，請參閱我們的自由線程 Docker 環境。不妨在“issues”部分中發布您的問題，并在您的用例中試用自由線程 Python！