检查 Colab 中的 GPU

如果你是个 GPU 程序员，你可能需要在 Colab 中使用 GPU。这篇文章介绍了如何检查 Colab 中的 GPU。在这篇文章中，我们将涵盖以下主题：

• 检查 Colab 运行时对象
• 安装和使用 nvidia-smi 命令
• 利用 torch.cuda 模块检测 GPU

检查 Colab 运行时对象

在打开 Colab 笔记本时，你需要从列表中选择一个运行时对象。你可以通过下面的代码块来检查你的运行时对象是否包含 GPU：

!nvidia-smi

如果你的运行时对象包含 GPU，你将看到类似下面的输出：

Tue May  4 15:30:23 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  A100-SXM4-40GB   WDDM    | 00000000:09:00.0 Off |                    0 |
| N/A   31C    P0    15W / 400W |      0MiB / 40536MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

安装和使用 nvidia-smi 命令

nvidia-smi 是一个 Nvidia 提供的监视 NVIDIA GPU 的命令行实用程序。你可以通过以下代码来检查你是否已经安装：

!command -v nvidia-smi

如果没有安装，你可以通过以下命令来安装：

!apt-get install -y nvidia-utils-450

作为替代，你也可以使用 Nvidia 提供的 Docker 映像，其中包含了 nvidia-smi 和关于 NVIDIA GPU 的其他工具。

!docker run --gpus all nvidia/cuda:11.0-base nvidia-smi

利用 torch.cuda 模块检测 GPU

如果你在使用 PyTorch 的话，你可以使用 torch.cuda 模块来检查你是否连接到了 GPU 上。下面的代码块演示了如何使用 torch.cuda 模块进行 GPU 检查：

import torch

if torch.cuda.is_available():
    device_id = torch.cuda.current_device()
    gpu_properties = torch.cuda.get_device_properties(device_id)
    print(f'Found GPU {gpu_properties.name} with max memory {gpu_properties.total_memory}bytes.')
else:
    print('No GPU available.')

如果你连接到 GPU 上，你将会看到类似下面的输出：

Found GPU Tesla T4 with max memory 15812263936 bytes.

如果你没有连接到 GPU 上，你将会看到以下输出：

No GPU available.