问题:
pytorch默认使用第一块显卡 cuda o ,在运行bert代码时,因为长文本,我把padsize调到512。导致显存不足,会占用大量的显存。程序会反复报Runtime Error,cuda out of memory的错。
分析:
服务器有三块显卡0,1,2。为啥空间不足,原来pytorch在初始化的时候会默认在第0块显卡上进行,这就导致,在第0块显卡空闲内存不多时,反复报错
解决:
通过几行代码解决问题
开头添加:
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2"
主函数模型训练部分使用并行:
# train
model = x.Model(config).to(config.device)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model,device_ids = [0, 1, 2])
train(config, model, train_iter, dev_iter, test_iter)
使用nn.DataParallel函数来用多个GPU来加速训练。
代码成功运行。。。。。。
再来看一下显存效果:
watch -n 0.2 nvidia-smi
看起来都在跑