解决cuda out of memory

问题:

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

 看起来都在跑

  • 3
    点赞
  • 19
    收藏
    觉得还不错? 一键收藏
  • 1
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值