如何选择在哪个 GPU 上运行作业

在深度学习中，我们通常需要利用 GPU 进行模型训练。对于一台有多个 GPU 的服务器来说，我们可以选择在哪个 GPU 上运行作业。那么如何选择 GPU 呢？

1. 查看服务器 GPU 信息

首先，我们需要查看服务器的 GPU 信息。可以使用 nvidia-smi 命令来查看服务器上的 GPU 状态和使用情况。

nvidia-smi

2. 利用环境变量选择 GPU

我们可以使用环境变量 CUDA_VISIBLE_DEVICES 来控制使用哪些 GPU。CUDA_VISIBLE_DEVICES 可以设置为以下值：

• 使用单个 GPU：设置为 GPU 的编号，从 0 开始计数。例如，使用第一个 GPU，可以设置为 CUDA_VISIBLE_DEVICES=0。
• 使用多个 GPU：设置为 GPU 编号的列表，以逗号分隔。例如，使用第一个和第二个 GPU，可以设置为CUDA_VISIBLE_DEVICES=0,1。

在执行命令时，可以在前面添加 CUDA_VISIBLE_DEVICES 环境变量设置。例如：

CUDA_VISIBLE_DEVICES=0 python train.py

这样就可在第一个 GPU 上运行训练。

3. 使用 Python 代码选择 GPU

除了使用环境变量外，我们还可以在 Python 代码中选择 GPU。可以使用 torch.cuda.set_device 方法来选择 GPU。例如：

import torch

torch.cuda.set_device(0)  # 选择第一个 GPU

总结

在选择使用哪个 GPU 运行作业时，我们可以使用环境变量 CUDA_VISIBLE_DEVICES 或者使用 Python 代码进行选择。无论哪种方法，掌握好 GPU 选择的技巧都有利于更高效地使用服务器资源。