我在使用DataParallel進(jìn)行雙GPU訓(xùn)練一個(gè)分類模型時(shí),定義batch size=16,然后遇到錯(cuò)誤:計(jì)算CrossEntropyLoss時(shí)predict與target的batch維度不一致,target的batch維度仍然是16,而模型輸出的predict的batch維度竟然等于32!
按照我以前對(duì)pytorch并行化的理解,一個(gè)batch的數(shù)據(jù)應(yīng)該是被平均劃分到不同GPU上并行計(jì)算,然后再將結(jié)果匯總的,這樣的話每個(gè)GPU上計(jì)算的batch維度大小=8,最后匯總batch size應(yīng)該仍然等于定義的16。
查閱DataParallel文檔,看到這么一句話:
Arbitrary positional and keyword inputs are allowed to be passed into DataParallel but some types are specially handled. tensors will be scattered on dim specified (default 0). tuple, list and dict types will be shallow copied. The other types will be shared among different threads and can be corrupted if written to in the model’s forward pass.
注意這句話:
tuple, list and dict types will be shallow copied.
然后再看我模型forward函數(shù)的輸入是一個(gè)list。問題就出在這兒,兩張卡并行計(jì)算時(shí)pytorch并不是將輸入的list劃分然后分到不同GPU中,而是直接淺拷貝后分發(fā),這就造成同一份數(shù)據(jù)被復(fù)制n份分發(fā)到n個(gè)GPU中,而不是我預(yù)想的一份數(shù)據(jù)被平均劃分為n份!
再去檢查模型輸出的tensor,果然后16個(gè)元素的值與前16個(gè)元素的值是一樣的。
因此,通過這件事,我們可以總結(jié)得到:
Pytorch中當(dāng)用DataParallel或DistributedDataParallel來進(jìn)行并行化計(jì)算時(shí),模型的輸入數(shù)據(jù)應(yīng)該要是一個(gè)tensor的類型。