设置 CUDA 可见设备

当使用 CUDA 加速训练深度学习模型时，为了充分利用 GPU 资源，通常需要设置 CUDA 可见设备，使程序只在指定的 GPU 上运行。

检查 CUDA 环境

在进行 CUDA 可见设备的设置前，需要先检查当前环境是否支持 CUDA。

import torch

if torch.cuda.is_available():
    print('CUDA is available!')
else:
    print('CUDA is not available.')

如果输出结果为 CUDA is available!，则说明当前环境已支持 CUDA 加速。否则，需要安装 CUDA 软件包并配置 CUDA 环境。

查看可用 GPU

一般来说，一台机器上会安装多个 GPU，我们需要先查看可用的 GPU，并确定要使用哪个 GPU 进行训练。可以通过 nvidia-smi 命令或 PyTorch 提供的函数 torch.cuda.device_count() 来查看。

import os

os.system('nvidia-smi')  # 查看当前可用 GPU

print(torch.cuda.device_count())  # 查看当前可用的 GPU 数量

设置 CUDA 可见设备

在确定要使用的 GPU 后，我们需要将其设置为 CUDA 可见设备，使得程序只在指定的 GPU 上运行，避免浪费其他 GPU 的资源。

import os
import torch

os.environ["CUDA_VISIBLE_DEVICES"] = "2"  # 设置 CUDA 可见设备为第 3 个 GPU

if torch.cuda.is_available():
    device = torch.device("cuda")
    print(f'Current device: {torch.cuda.current_device()}')
else:
    device = torch.device("cpu")
    print('You should install CUDA and configure environment.')

上述代码将 CUDA 可见设备设置为第 3 个 GPU，即 ID 为 2 的 GPU。我们可以通过 torch.cuda.current_device() 函数来查看当前使用的 GPU。

如果只有一张 GPU 可用，那么可以直接将 CUDA 可见设备设置为 0，也可以不设置，直接使用默认的第 0 张 GPU。

总结

通过以上介绍，我们可以轻松地设置 CUDA 可见设备，避免浪费 GPU 资源，充分利用硬件加速深度学习模型的训练。