【Q&A】pytorch中的worker如何工作的

目錄

• Question
• Answer
• 參考資料

博客地址： https://www.cnblogs.com/hesse-summer/，歡迎轉載，轉載請說明出處。

Question

一直很迷，
在給Dataloader設置worker數量（num_worker）時，到底設置多少合適？這個worker到底怎么工作的？
如果將num_worker設為0（也是默認值），就沒有worker了嗎？

worker的使用場景：

from torch.utils.data import DataLoader

train_loader = DataLoader(dataset=train_data, batch_size=train_bs, shuffle=True, num_worker=4)

valid_loader = DataLoader(dataset=valid_data, batch_size=valid_bs, num_worker=4)

Answer

1. 每每輪到dataloader加載數據時：

   for epoch in range(start_epoch, end_epoch):
       for i, data in enumerate(trainloader):
   

   dataloader一次性創建num_worker個worker，（也可以說dataloader一次性創建num_worker個工作進程，worker也是普通的工作進程），

   並用batch_sampler將指定batch分配給指定worker，worker將它負責的batch加載進RAM。

   然后，dataloader從RAM中找本輪迭代要用的batch，如果找到了，就使用。如果沒找到，就要num_worker個worker繼續加載batch到內存，直到dataloader在RAM中找到目標batch。一般情況下都是能找到的，因為batch_sampler指定batch時當然優先指定本輪要用的batch。

2. num_worker設置得大，好處是尋batch速度快，因為下一輪迭代的batch很可能在上一輪/上上一輪...迭代時已經加載好了。壞處是內存開銷大，也加重了CPU負擔（worker加載數據到RAM的進程是CPU復制的嘛）。num_workers的經驗設置值是自己電腦/服務器的CPU核心數，如果CPU很強、RAM也很充足，就可以設置得更大些。

3. 如果num_worker設為0，意味着每一輪迭代時，dataloader不再有自主加載數據到RAM這一步驟（因為沒有worker了），而是在RAM中找batch，找不到時再加載相應的batch。缺點當然是速度更慢。

參考資料

1. https://stackoverflow.com/questions/53998282/how-does-the-number-of-workers-parameter-in-pytorch-dataloader-actually-work
2. https://discuss.pytorch.org/t/guidelines-for-assigning-num-workers-to-dataloader/813