1. 检查GPU使用情况

由于nvidia-smi命令显示的gpu id和实际运行时gpu id并不对应，因此通过上述代码来查看实际运行时的gpu id

import torch

def print_gpu_info():
    # 检查可用的 GPU 数量
    num_gpus = torch.cuda.device_count()

    if num_gpus > 0:
        print(f"可用的 GPU 数量: {num_gpus}")

        # 打印每个 GPU 设备的 ID 和名称
        for device_id in range(num_gpus):
            device_name = torch.cuda.get_device_name(device_id)
            print(f"GPU 设备 {device_id}: {device_name}")
    else:
        print("没有可用的 GPU")

# 调用函数打印 GPU 信息
print_gpu_info()

# 设置要使用的 GPU 设备 ID
device_id = 0
torch.cuda.set_device(device_id)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device:", device)

# 获取当前设备的 ID 和名称
device_id = torch.cuda.current_device()
device_name = torch.cuda.get_device_name(device_id)
# 打印使用的 GPU 设备
print(f"当前使用的 GPU 设备: {device_name}")
print('GPU id: ' + str(device_id))

2. 设置使用特定的GPU

在PyTorch中，你可以通过设置 CUDA_VISIBLE_DEVICES 环境变量来控制哪些GPU对当前进程可见，从而限制PyTorch使用的GPU数量。这个环境变量接受一个由逗号分隔的GPU索引列表。

如果你只想让PyTorch使用特定的几个GPU，可以在运行Python脚本或程序之前设置这个环境变量。

在命令行中，你可以这样设置（假设你想使用GPU 0和GPU 2）：

export CUDA_VISIBLE_DEVICES=0,2

然后你可以运行你的PyTorch程序，程序将只看到你指定的GPU。

在Python脚本中，你可以使用 os.environ 在程序开始时设置：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,2"

import torch
# 现在PyTorch只能看到GPU 0和GPU 2

请注意，当你设置 CUDA_VISIBLE_DEVICES 后，PyTorch看到的GPU索引将会重新排序以匹配这个变量。也就是说，即使你设置 CUDA_VISIBLE_DEVICES=2,3，在PyTorch中这两个设备仍然会被标记为 cuda:0 和 cuda:1。