pytorch 解決out of memory

主要是針對指定了兩片GPU，但默認使用第一塊的問題。

我在用pytorch跑程序的時候，由於代碼比較大，需要用兩塊板子。但無論是在代碼中加入

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

還是 直接在命令行輸入

CUDA_VISIBLE_DEVICES='0,1' python train.py

都無法解決問題。後來發現應使用 torch.nn.DataParallel這個函數,具體用法如下所示：

#自動選擇GPU
model = nn.DataParallel(model)
model = model.cuda()
#選擇指定GPU
model = torch.nn.DataParallel(model.cuda(),device_ids=[2,3])
#指定第二、三塊板子

但是發現報錯：AssertionError: Invalid device id，這是因爲pytorch默認使用gpu編號爲device:0的設備，對模型指定編號gpu訓練時，必須要有編號爲device:0的gpu，不然會報AssertionError: Invalid device id錯誤；所以如果你想使用2、3塊板子，需要修改pytorch默認編號，故代碼修改爲：

torch.cuda.set_device(2)#指定默認編號爲2
model = torch.nn.DataParallel(model.cuda(),device_ids=[2,3])

還有一種方法是：

import os 
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
model = torch.nn.DataParallel(model.cuda(),device_ids=[2,3])
#pytorch 會默認把device：2改爲device：0，而device：3變爲device：1

具體可參考這篇博客.
問題解決！歐耶！