指定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)

相关推荐
贺一航【Niki】5 小时前
【学习笔记】杂乱知识
笔记·学习
つ安静与叛逆的小籹人6 小时前
小红书API:通过笔记ID获取笔记详情数据教程
笔记·python
ClutchoQ8 小时前
【你指的API是哪个API?软件工程师跨服聊天实录】
笔记·其他
二哈赛车手10 小时前
新人笔记---Spring AI的Advisor以及其底层机制讲解(涉及源码),包含一些遇见的Spring AI的Advisor缺陷问题的解决方案
java·人工智能·spring boot·笔记·spring
郝学胜-神的一滴12 小时前
反向传播:神经网络的「灵魂」修炼法则
人工智能·pytorch·深度学习·神经网络·机器学习·数据挖掘
Tutankaaa12 小时前
知识竞赛软件SaaS版 vs 本地部署
人工智能·经验分享·笔记·学习
许长安12 小时前
RPC 异步调用基本使用方法:基于官方helloworld-async 示例
c++·经验分享·笔记·rpc
sheeta199814 小时前
苍穹外卖Day12笔记
笔记
hnult14 小时前
AI 在线考试全流程 :考试云七大核心 AI 功能深度解读
人工智能·笔记·重构·课程设计
普通young man14 小时前
vim常用编辑和视图(个人笔记)
笔记·编辑器·vim