dataset = SimpleDataset() dataloader = DataLoader(dataset, num_workers=2, batch_size=3) for batch in dataloader: print(batch)
主進程初始化
dataloader = DataLoader(dataset, num_workers=2, batch_size=3)
- 創建
num_workers
個不同的子進程。 - 為每個子進程維護一個
index_queue
,用於發送需要采集數據的index。注意,隊列每個元素都是一個batch的index的數組。 - 所有子進程共享一個
data_queue
,用於存放子進程返回的數據。 - 為每個子進程的
index_queue
依次存放兩個batch的數據。比如本代碼,會為0號worker存放[(0,1,2),(6,7,8)]
,為1號worker存放[(3,4,5),(9,10,11)]
。依次存放是為了盡量保證取得數據的順序和采樣器生成的順序是一致的。而存放兩組數據是為了實現數據預取的功能(即取第i個batch的數據時,i+1個batch的數據已經讀好,並開始緩存i+2個batch的數據。
基於上述分析,我們可以很容易得到以下幾個工程經驗:
- 一個完整的batch由一個worker產生,而不是由多個worker共同產生。這樣做主要是為了盡可能減少進程通信的次數。
- 內存占用有多大?正常情況下,內存中穩定存儲
2*num_workers
個batch的數據。 但由於亂序緩存機制的存在,實際會比這個大一些。最壞情況是某個worker卡住始終不返回數據,那就會一直緩存后續數據,一直讀取,直到爆內存。個人認為這個設計並不好,為何不每個worker維護自己的data_queue? - 數據預取的trick還需要加嗎?之前在很多文章看到自己定義DataLoader,然后手動prefetch下一個batch的數據以加速。個人認為由於雙緩存機制,這項trick已經沒用了。
- 在設計
Dataset
時,應該讓每次__getitem__
的時間盡可能均勻,這樣不會浪費過多的內存在亂序緩存機制上,可以有效降低內存占用。
問題:
在跑tsm模型的時候,發現gpu的利用率低,而且中間有間斷,num_workers設置的是16,batch_size16,有8個顯卡(后來檢查才發現有56個cpu)
然后把num_workers改成56,batchsize也改為56,解決!!!gpu利用率一下子提升了