pytorch: n個(gè)GPU并行計(jì)算時(shí)模型輸出的batch size等于預(yù)定義bs的n倍

我在使用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的類型。

最后編輯于
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時(shí)請(qǐng)結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點(diǎn),簡(jiǎn)書系信息發(fā)布平臺(tái),僅提供信息存儲(chǔ)服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容