指定显卡的三种方式

EnochChen_2026-01-30 13:32

第一种 os.environ["CUDA_VISIBLE_DEVICES"] = "1"

python 复制代码

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import torch

print(torch.cuda.device_count())
print(torch.cuda.current_device())
print(torch.cuda.get_device_name(0))

结果

物理 GPU 1 被映射为逻辑 GPU 0

python 复制代码

1
0
NVIDIA GeForce RTX 4090

第二种命令行 CUDA_VISIBLE_DEVICES=1 python cuda_test.py

python 复制代码

import torch

print(torch.cuda.device_count())
print(torch.cuda.current_device())
print(torch.cuda.get_device_name(0))

结果

物理 GPU 1 被映射为逻辑 GPU 0

python 复制代码

1
0
NVIDIA GeForce RTX 4090

第三种 torch.cuda.set_device(1)

python 复制代码

import torch
torch.cuda.set_device(1)

print(torch.cuda.device_count())
print(torch.cuda.current_device())
print(torch.cuda.get_device_name(0))
print(torch.cuda.get_device_name(1))

结果

设置任意cuda：1设备为当前cuda设备

python 复制代码

2
1
NVIDIA GeForce RTX 4090
NVIDIA GeForce RTX 4090

指定显卡的三种方式

第一种 os.environ["CUDA_VISIBLE_DEVICES"] = "1"

第二种 命令行 CUDA_VISIBLE_DEVICES=1 python cuda_test.py

第三种 torch.cuda.set_device(1)

第二种命令行 CUDA_VISIBLE_DEVICES=1 python cuda_test.py