
它們只分擔「CPU 端的資料前處理」,完全不分擔「GPU 端的模型運算」。 forward / backward / 梯度更新永遠是主行程一個人在 GPU 上做,4 個 worker 做的是讀圖、做資料增強、把標註轉成 tensor 這些跟模型本身無關、但訓練前必須做的雜活。
為什麼要開「行程」而不是「執行緒」
讀圖、resize、mosaic/隨機增強這些工作是 CPU + 硬碟 I/O bound。如果用 Python 的 threading 來做,會被 GIL(全域解釋器鎖)卡住——同一時間還是只有一個執行緒在真正執行 Python bytecode,多執行緒在這種 CPU 密集工作上幾乎沒有平行效益。PyTorch 的 DataLoader 因此選擇用 multiprocessing:每個 worker 是一個完全獨立的作業系統行程,各自有自己的 Python 解釋器、自己的 GIL,才能真正利用到多核心 CPU 同時處理。
這也正是上次那個 bug 的根源:它們是獨立的 OS process,不是執行緒,所以主行程被 TerminateProcess 砍掉時,作業系統不會自動連帶砍掉它們。
任務怎麼分配——以 batch 為單位、輪流(round-robin)
主行程裡有一個 index sampler,會把資料集的索引依照 batch 大小切成一份份的 index 清單,然後輪流塞進每個 worker 各自的 index queue:batch 0 給 worker 0、batch 1 給 worker 1、batch 2 給 worker 2、batch 3 給 worker 3、batch 4 再回到 worker 0……以此類推。每個 worker 收到自己的 index 清單後,完全獨立地呼叫 dataset.__getitem__:讀圖(cv2.imread)、跑增強流程(letterbox/resize、mosaic、翻轉等)、解析對應的 .txt/.json 標註、組成 tensor,做完就丟進一個大家共用的 result queue。
主行程的訓練迴圈則是從 result queue 裡按照原本的 batch 順序把結果取出來(即使 worker 完成的順序可能因為速度不同而錯開,主行程會用內部的索引機制重新排好順序再交給模型),丟到 GPU 上跑這一個 batch 的 forward/backward。
Prefetch:讓 GPU 不會空等
worker 不是做完一批就停下來等主行程要,預設會提前準備好幾個 batch(prefetch_factor,預設 2)排在 queue 裡。這樣主行程的 GPU 在算這一個 batch 的同時,CPU 這邊 4 個 worker 已經在背景把下幾個 batch 準備好了——這才是「workers > 0」真正帶來效益的地方:讓 CPU 前處理跟 GPU 運算重疊(overlap),而不是讓 GPU 算完一批就乾等資料。
Windows 上為什麼特別容易變孤兒——回到上次的 bug
這裡有個跟 Linux 不一樣的關鍵差異:Linux 可以用 fork() 複製主行程的記憶體來開 worker,速度快也不用重新 import;但 Windows 沒有 fork(),只能用 spawn——每個 worker 其實是重新啟動一個全新的 python.exe,重新 import 一次你整支 script(包括 from ultralytics import YOLO),再把 dataset 物件 pickle 過去重建。這就是為什麼 yolo_runner.py 產生的腳本一定要寫成:
if __name__ == '__main__':
multiprocessing.freeze_support()
main()
少了這層保護,Windows 上 spawn 出來的子行程會把整支腳本當成新的進入點再執行一次,遞迴炸出無窮多個訓練行程。
正常情況下,當主行程正常結束(跑完、Ctrl+C 等乾淨退出路徑)時,PyTorch 會透過 atexit 鉤子主動通知所有 worker 關閉。但 QProcess.kill() 在 Windows 上等於直接呼叫 TerminateProcess——是粗暴地把整個行程瞬間砍掉,完全沒有機會跑到那些 atexit 清理邏輯。於是 4 個 worker 行程瞬間失去了會跟它們對話的父行程,變成孤兒,繼續占著 CPU(甚至可能還握著 GPU context),這正是會拖慢你之前遇到的「正在驗證資料集…」那個純 CPU 迴圈的元兇——也是我們改用 taskkill /F /T /PID 整條行程樹一次清掉的原因。