【PyTorch】设置CUDA_VISIBLE_DEVICES无效的问题以及多卡使用以及CUDA out of memory问题

方法1:

理想情况下,该环境变量应设置在程序的顶部。如果在设置 torch.backends.cudnn.benchmark 之后调用 CUDA_VISIBLE_DEVICES 变量,则更改 CUDA_VISIBLE_DEVICES 变量将不起作用。

复制代码
import os
os.environ["CUDA_VISIBLE_DEVICES"]=<YOUR_GPU_NUMBER_HERE>

方法2:

shell 复制代码
$ CUDA_VISIBLE_DEVICES=1,2 python train.py

REF


问题:单机多卡使用时,指定device不是从0开始时,出现AssertionError: Invalid device id错误

原因:

pytorch默认使用gpu编号为device:0的设备,可以使用

torch.nn.DataParallel(model, device_ids=0, 1)

对模型进行制定gpu指定编号多gpu训练,必须要有编号为device:0的gpu,不然会报AssertionError: Invalid device id错误;

如果当gpu编号为device:0的设备被占用时,指定其他编号gpu使用torch.nn.DataParallel(model, device_ids=1, 2)指定gpu编号会出现AssertionError: Invalid device id错误,这是因为pytorch默认使用gpu编号为device:0的设备,需要修改pytorch默认编号;

解决方法

方法1:

python 复制代码
# 使用set_device函数来指定默认GPU使用编号
torch.cuda.set_device(6)

#使用方式
torch.nn.DataParallel(model,device_ids=[6,7,8])

缺点:

  • 必须包含set_device(1)指定的device:1的设备;
  • 且仍然会存在占用一些device:0的gpu内存;

方法2:(推荐)

python 复制代码
# 使用如下方式来指定使用设备,这样会会把device:6改为device:0,
# device:7改为device:1,以此类推. 则pytorch默认的编号还是以device:0开始.
os.environ["CUDA_VISIBLE_DEVICES"] = "6,7,8"

# 指定方式
torch.nn.DataParallel(model,device_ids=[0,1,2])

REF


**问题:**单机多卡训练时,代码只运行了gpu0,并导致CUDA out of memory. 如下图所示:
RuntimeError: CUDA out of memory. Tried to allocate 5.18 GiB (GPU 0; 31.74 GiB total capacity; 21.11 GiB already allocated; 5.10 GiB free; 25.60 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.

解决方法:

with torch.cuda.device(7):

opt = opts().parse()

main(opt)

pytorch在初始化的时候会默认在第0块显卡上进行,会占用一定的显存,这就导致,在第0块显卡空闲内存不多时,程序会反复报上面的错误。

torch.nn.DataParallel

pytorch单机多卡最简单的实现方法就是使用nn.DataParallel类,其几乎仅使用一行代码net = torch.nn.DataParallel(net)就可让模型同时在多张GPU上训练,它大致的工作过程如下图所示:

DataParallel前传与反传工作过程

在每一个Iteration的Forward过程中,nn.DataParallel都自动将输入按照gpu_batch进行split,然后复制模型参数到各个GPU上,分别进行前传后将得到网络输出,最后将结果concat到一起送往0号卡中。

**在Backward过程中,先由0号卡计算loss函数,**通过loss.backward()得到损失函数相于各个gpu输出结果的梯度grad_l1 ... gradln,接下来0号卡将所有的grad_l送回对应的GPU中,然后GPU们分别进行backward得到各个GPU上面的模型参数梯度值gradm1 ... gradmn,最后所有参数的梯度汇总到GPU0卡进行update。

注:DataParallel的整个并行训练过程利用python多线程实现

由以上工作过程分析可知,nn.DataParallel有着这样几个无法避免的问题:

  • 负载不均衡问题。gpu0所承担的任务明显要重于其他gpu
  • 速度问题。每个iteration都需要复制模型且均从GPU0卡向其他GPU复制,通讯任务重且效率低;python多线程GIL锁导致的线程颠簸(thrashing)问题。
  • 只能单机运行。由于单进程的约束导致。
  • 只能切分batch到多GPU,而无法让一个model分布在多个GPU上。当一个模型过大,设置batchsize=1时其显存占用仍然大于单张显卡显存,此时就无法使用DataParallel类进行训练。
    因此官方推荐使用torch.nn.DistributedDataParallel替代nn.DataParallel,即使是单机多卡的情况下。

REF

相关推荐
“AI国潮设计-小江”7 小时前
【Python/SDXL实战】潮汕国潮IP视觉落地:普宁英歌舞猫IP & 创意甜品设计(附ComfyUI工作流与商业授权说明)
开发语言·人工智能·python·prompt·aigc
智慧大脑搬运工7 小时前
美丽蓝天政策申报|从技术目录到金融对接:环保科技成果转化的政策链路解析
人工智能
AI 思录7 小时前
AI生造词、乱用术语怎么治?一份防“词汇污染“的Prompt约束模板
人工智能·语言模型·prompt·ai写作·用户体验·ai合规
529宝宝起名网7 小时前
用 Python 分析汉字字源与起名用字的关联规律:从六书结构到文化寓意的起名偏好洞察
开发语言·python
hsg777 小时前
简述:人工智能 + 软件实施方案
人工智能
一只旭宝7 小时前
Python 与 C/C++ 内存模型对比总结
c语言·c++·python
人工智能时代 准备好了吗7 小时前
AI写错品牌一个字,如何判断是名称误差还是认错对象?
人工智能
电子制造自留地8 小时前
AI服务器PCB的超低损耗材料选型逻辑
运维·服务器·人工智能·科技·制造·科普·pcb工艺
张彦峰ZYF8 小时前
从“全量 OCR”到按页智能路由:pdf-inspector 与企业 PDF 解析架构的工程化重构
人工智能·pdf·ocr·ai agent·pdf-inspector
ocean21038 小时前
2025-2026年AI应用开发与Agent面试高频知识点洞察
人工智能·面试·职场和发展