指定GPU设备

在运行代码的时候遇到下面报错:

完整报错信息:

python 复制代码
DeferredCudaCallError: CUDA call failed lazily at initialization with error: device >= 0 && device < num_gpus INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1716905969824/work/aten/src/ATen/cuda/CUDAContext.cpp":50, please report a bug to PyTorch. device=, num_gpus=

排查后发现问题出在有1块GPU error了,在代码中显式地指定GPU编号即可解决问题。

指定GPU设备

方法一:

python 复制代码
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2'

CUDA_VISIBLE_DEVICES :决定进程能看见哪些物理 GPU,并且会重新编号(映射成从 0 开始的逻辑编号),要放在import torch前。

方法二:

python 复制代码
import torch
torch.cuda.set_device(2)

torch.cuda.set_device(k) 是在"已经可见的逻辑 GPU 列表"里,把默认 GPU 指针设成第 k 个;不会屏蔽 GPU,也不会改变可见性。

简单来说:

  • 想"只用某几张卡 / 不要碰其他卡" → 用CUDA_VISIBLE_DEVICES来设置;
  • 想"在可见卡里把默认卡设成哪张" → 用 set_device

示例1:

python 复制代码
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
import torch
torch.cuda.set_device(1)

这里其实是在选:逻辑1->GPU 3,而不是GPU 1。

示例2:

python 复制代码
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
import torch
torch.cuda.set_device(2)   # 报错

因为可见逻辑设备只有0和1,没有2。

参考:

1\] [【设置gpu设备】os.environ\['CUDA_VISIBLE_DEVICES'\] 和 torch.cuda.set_device()_os设置gpu使用设备](https://blog.csdn.net/qq_51392112/article/details/129739768)

相关推荐
西岸行者4 天前
学习笔记:SKILLS 能帮助更好的vibe coding
笔记·学习
starlaky4 天前
Django入门笔记
笔记·django
勇气要爆发4 天前
吴恩达《LangChain LLM 应用开发精读笔记》1-Introduction_介绍
笔记·langchain·吴恩达
悠哉悠哉愿意4 天前
【单片机学习笔记】串口、超声波、NE555的同时使用
笔记·单片机·学习
勇气要爆发4 天前
吴恩达《LangChain LLM 应用开发精读笔记》2-Models, Prompts and Parsers 模型、提示和解析器
android·笔记·langchain
qianshanxue114 天前
计算机操作的一些笔记标题
笔记
土拨鼠烧电路4 天前
笔记11:数据中台:不是数据仓库,是业务能力复用的引擎
数据仓库·笔记
土拨鼠烧电路4 天前
笔记14:集成与架构:连接孤岛,构建敏捷响应能力
笔记·架构
烟花落o4 天前
栈和队列的知识点及代码
开发语言·数据结构·笔记·栈和队列·编程学习