yolov5訓練時引數workers與batch-size的深入理解

2022-03-24 19:00:47

yolov5訓練命令

 python .train.py --data my.yaml --workers 8 --batch-size 32 --epochs 100

yolov5的訓練很簡單，下載好倉庫，裝好依賴後，只需自定義一下data目錄中的yaml檔案就可以了。這裡我使用自定義的my.yaml檔案，裡面就是定義資料集位置和訓練種類數和名字。

一般訓練主要需要調整的引數是這兩個：

指資料裝載時cpu所使用的執行緒數，預設為8。程式碼解釋如下

parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')

一般默使用8的話，會報錯~~。原因是爆系統記憶體，除了實體記憶體外，需要調整系統的虛擬記憶體。訓練時主要看已提交哪裡的實際值是否會超過最大值，超過了不是強退程式就是報錯。

所以需要根據實際情況分配系統虛擬記憶體(python執行程式所在的盤)的最大值

就是一次往GPU哪裡塞多少張圖片了。決定了視訊記憶體佔用大小，預設是16。

parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')

訓練時視訊記憶體佔用越大當然效果越好，但如果爆視訊記憶體，也是會無法訓練的。我使用–batch-size 32時，視訊記憶體差不多能利用完。

對於workers，並不是越大越好，太大時gpu其實處理不過來，訓練速度一樣，但虛擬記憶體（磁碟空間）會成倍佔用。

workers為4時的記憶體佔用

workers為8時的記憶體佔用

我的顯示卡是rtx3050，實際使用中上到4以上就差別不大了，gpu完全吃滿了。但是如果設定得太小，gpu會跑不滿。比如當workers=1時，顯示卡功耗只得72W，速度慢了一半；workers=4時，顯示卡功耗能上到120+w，完全榨乾了顯示卡的算力。所以需要根據你實際的算力調整這個引數。

2. 對於batch-size，有點玄學。理論是能儘量跑滿視訊記憶體為佳，但實際測試下來，發現當為8的倍數時效率更高一點。就是32時的訓練效率會比34的高一點，這裡就不太清楚原理是什麼了，實際操作下來是這樣。

以上引數的調整能最大化顯示卡的使用效率，其中的具體數值和電腦的實際設定還有模型大小、資料集大小有關，需要根據實際情況反覆調整。當然，要實質提升訓練速度，還是得有好顯示卡（鈔能力）~~~~