admin 管理员组

文章数量: 1184232

问题：

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

 看起来都在跑

本文标签： CUDA memory

版权声明：本文标题：解决cuda out of memory 内容由网友自发贡献，该文观点仅代表作者本人， 转载请联系作者并注明出处：http://www.roclinux.cn/b/1766204200a3444332.html， 本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如发现本站有涉嫌抄袭侵权/违法违规的内容，一经查实，本站将立刻删除。