那些年在pytorch上踩過的坑

今天又發(fā)現(xiàn)了一個pytorch的小坑，給大家分享一下。手上兩份同一模型的代碼，一份用tensorflow寫的，另一份是我拿pytorch寫的，模型架構(gòu)一模一樣，預(yù)處理數(shù)據(jù)的邏輯也一模一樣，測試發(fā)現(xiàn)模型推理的速度也差不多。一份預(yù)處理代碼是為pytorch模型寫的，用到的庫是torch，另一份是為tensorflow寫的，用到的是numpy。在訓(xùn)練時，每個epoch耗時居然差距非常大，pytorch的代碼在140w條數(shù)據(jù)上訓(xùn)練每輪耗時約45min，而tensorflow版的代碼耗時僅約12min。

我把代碼看了又看，百思不得其解，預(yù)處理的代碼比較復(fù)雜，都包含兩個for循環(huán)，pytorch版代碼我把更多的預(yù)處理步驟放到了Dataset里，這樣訓(xùn)練時加載每個batch后，再要處理的步驟就更少了，速度也應(yīng)該更快，而tensorflow版代碼的for循環(huán)里預(yù)處理的步驟明明更多，怎么會速度比我的代碼還快呢？然而，經(jīng)過我的測試發(fā)現(xiàn)，從加載每個batch的數(shù)據(jù)進來開始，經(jīng)過預(yù)處理，直到輸入到模型做計算前，兩者的耗時差了約7~8倍。最后發(fā)現(xiàn)問題出在對pytorch的tensor進行了頻繁的索引操作。

下面做個實驗給大家直觀體驗一下，對tensor做索引和對array做索引的速度差距有多大，tensor和array都是大小(1000x1000)的二維數(shù)組。

Pytorch(version==1.4.1)索引1000000次耗時：3.51秒

Numpy索引1000000次耗時：0.43秒

我還特意對比了一下對TensorFlow的tensor做索引的耗時

TensorFlow(version==2.1.0)索引1000000次耗時：118.89秒

由此可見tensor和array的索引速度至少差距在10倍，不過這也在情理之中，畢竟tensor要比array“重”得多。因此在使用pytorch和tensorflow時，頻繁需要索引的操作一定要先把tensor轉(zhuǎn)換為numpy.array來做！

除此之外，與其對二維數(shù)組進行索引，不如將其展平為一維數(shù)組，算上展平的時間，速度還會有不少提升。

Pytorch從3.51秒降到了1.94秒

Numpy從0.43秒降到了0.29秒

如果在訓(xùn)練和數(shù)據(jù)預(yù)處理過程中發(fā)現(xiàn)自己的代碼跑起來速度非常慢，記得看一看有沒有對tensor做太多次索引，如果有的話，要把它轉(zhuǎn)為numpy.array，還有，盡量把二維、三維的索引變成一維的索引，這些都能加快你訓(xùn)練模型的速度。

PS：最后我的代碼終于訓(xùn)練一輪也只需要不到12min了，后來又找了點加速的辦法，把訓(xùn)練一輪的時間控制到了9min以內(nèi)，這些就放在以后再寫吧~