在进行深度学习或其他需要大量计算的任务时，使用 GPU（图形处理单元）可以显著提升计算性能。随着多款显卡的出现，如何查看、指定使用的 GPU 数量和编号，成为一个非常重要的问题。本篇文章将通过一些实用的代码示例，阐述如何在不同的深度学习框架中管理 GPU 的使用。
一、CUDA 和 cuDNN 安装
在使用 GPU 进行深度学习之前，请确保你已经安装了 NVIDIA 的 CUDA 工具包和 cuDNN。这两个组件能够帮助深度学习框架有效利用 GPU 进行加速计算。
二、查看系统中的 GPU
在使用之前，首先需要确认系统中装备了哪些 GPU。可以使用以下命令：
nvidia-smi
该命令会显示当前 GPU 的数量、型号、显存使用情况等信息。例如：
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 465.19.01 Driver Version: 465.19.01 CUDA Version: 11.3 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| GPU-0: GeForce GTX 1080 Ti Off | 00000000:01:00.0 Off | N/A |
| 30% 50C P8 10W / 250W | 0MiB / 11178MiB | 0% Default |
+-------------------------------+----------------------+----------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID ID Type Usage N/A Usage |
| 0 N/A N/A N/A N/A N/A N/A |
+-----------------------------------------------------------------------------+
三、在 PyTorch 中指定使用的 GPU
如果您使用的是 PyTorch 框架，可以通过设置设备来指定使用的 GPU。以下是一个示例代码：
import torch
# 查看可用的 GPU 数量
gpu_count = torch.cuda.device_count()
print(f"可用的 GPU 数量: {gpu_count}")
# 指定使用的 GPU 编号，比如使用第 0 张 GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 创建一个张量并将其放置在指定的 GPU 上
x = torch.randn(2, 3).to(device)
print(f"张量 x 在设备: {x.device}")
如果你想使用特定编号的多个 GPU，比如 GPU 0 和 GPU 1，可以使用 DataParallel:
model = MyModel()
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
# 将模型转移到 GPU
model.to(device)
四、在 TensorFlow 中指定使用的 GPU
在 TensorFlow 框架中，你同样可以通过简单的代码来指定使用的 GPU。例如：
import tensorflow as tf
# 查看可用的 GPU
gpus = tf.config.list_physical_devices('GPU')
print(f"可用的 GPU: {gpus}")
# 指定使用的 GPU
if gpus:
tf.config.set_visible_devices(gpus[0], 'GPU') # 只使用第一张 GPU
# 创建 TensorFlow 的模型
model = MyModel()
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
# 开始训练
model.fit(train_dataset, epochs=10)
五、总结
管理 GPU 的使用是深度学习过程中的一个重要环节。不同的框架提供了多种方式来查看和选择 GPU。在使用 GPU 进行高性能计算时，合理地选择和管理可以显著提升训练速度和效率。通过上述代码示例，相信你已经对如何查看和指定 GPU 有了一定的了解。在未来的深度学习项目中，希望你能灵活运用这些技巧，提升你的研究与工作的效率。
