起因

最近在修改上一個(gè)同事加載和預(yù)處理數(shù)據(jù)的代碼，原版的代碼使用tf1.4.1寫的，數(shù)據(jù)加載也是完全就是for循環(huán)讀取+預(yù)處理，每讀入并預(yù)處理好一個(gè)batch就返回丟給模型訓(xùn)練，如此往復(fù)，我覺得速度實(shí)在太慢了，而且我新寫的代碼都是基于pytorch，雖然預(yù)處理的過程很復(fù)雜，我還是下決心自己改寫。

用pytorch加載預(yù)處理數(shù)據(jù)，最常用的就是torch.utils.data.Dataset和torch.utils.data.DataLoader組合起來，把數(shù)據(jù)預(yù)處理都在Dataset里寫好，再在DataLoader里設(shè)定batch_size, shuffle等參數(shù)去加載數(shù)據(jù)，網(wǎng)上的教程非常多，這里我就不展開講了。

過程

現(xiàn)在我已經(jīng)獲得了train_loader和test_loader，可以從它們里面每次讀取一個(gè)batch出來訓(xùn)練，可照理說加載Dataset時(shí)占用了大量內(nèi)存是正常的，因?yàn)閿?shù)據(jù)都預(yù)加載好了，就只需要用DataLoader讀取就行了，但在訓(xùn)練的過程中，內(nèi)存不應(yīng)該隨著訓(xùn)練而逐漸增加。我眼睜睜看著內(nèi)存占用從8、9個(gè)g，逐漸漲到了25個(gè)g，程序最終因?yàn)檎紳M內(nèi)存而崩潰。檢查了半天自己的代碼都沒找出問題所在，后來用memory_profiler查看內(nèi)存占用情況，發(fā)現(xiàn)問題主要出現(xiàn)在這一行代碼：actual_labels += list(correctness)，correctness的類型是torch.FloatTensor，actual_labels是python原生的list。

pred和correctness是同類型同長度的tensor，在將其轉(zhuǎn)換為list再添加到已有的list中時(shí)，占用的內(nèi)存相差了約5個(gè)Mb，于是我將上面代碼改寫為：

問題解決了！

為什么下面的代碼就沒事，上面直接將Tensor轉(zhuǎn)為list就會(huì)發(fā)生這種奇怪的現(xiàn)象？我去github和知乎看到了遇到類似問題的issue與文章，下面是鏈接

1. https://github.com/pytorch/pytorch/issues/13246
2. https://zhuanlan.zhihu.com/p/86286137
3. https://github.com/pytorch/pytorch/issues/17499

結(jié)論

目前得到的結(jié)論大概是python list的design有問題，導(dǎo)致了這種情況發(fā)生，pytorch團(tuán)隊(duì)雖然竭力修復(fù)，但他們表示因?yàn)檫@是python設(shè)計(jì)的缺陷，超出了他們的能力范圍，上面第一個(gè)issue主要是針對(duì)DataLoader里num_workers>0時(shí)會(huì)導(dǎo)致內(nèi)存泄漏，里面也提到了list與tensor互轉(zhuǎn)亦會(huì)發(fā)生內(nèi)存泄漏，這個(gè)issue已經(jīng)一年多了還沒能close。

因此，在使用pytorch時(shí)，應(yīng)該盡力避免list的使用，一定不能讓tensor和list直接互相轉(zhuǎn)換，如果一定要做，應(yīng)該將tensor從cuda轉(zhuǎn)到cpu上，轉(zhuǎn)為numpy.array，最后轉(zhuǎn)為list，反之亦然。