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

具体可参考这篇博客.
问题解决！欧耶！