1.查看GPU

1.1.查看当前GPU是否可用

#cuda是否可用
torch.cuda.is_available()  # 返回 True 则可用, 返回 False 则不可用

1.2.查看GPU数量

# 返回gpu数量
torch.cuda.device_count()  # 返回值就是 GPU 数目

1.3.返回GPU名字

# 返回gpu名字，设备索引默认从0开始
torch.cuda.get_device_name(index)  # index 是索引, 默认从 0 开始

1.4.返回当前设备索引

# 返回当前设备索引
torch.cuda.current_device()

2.指定GPU

2.1.使用 CUDA_VISIBLE_DEVICES

2.1.1.在终端中指定

CUDA_VISIBLE_DEVICES=1 python main.py

2.1.2.代码中指定

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

上述两种方式其实都是添加环境变量，限制CUDA程序所能使用的GPU设备。从名字上，就能看出是限定程序能“看到”的设备。

CUDA应用运行时，CUDA将遍历当前可见的设备，并从零开始为可见设备编号。

可以同时指定多个设备（设备是真实存在的，不要指定不存在的设备）

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

2.2.使用函数 set_device

import torch
torch.cuda.set_device(id)

官方建议使用CUDA_VISIBLE_DEVICES，不建议使用 set_device 函数。