PyTorch 检查GPU版本是否安装成功的操作

当使用 PyTorch 进行深度学习计算的时候，GPU 可以大大加速模型训练的速度。因此，在安装 PyTorch 时，我们需要确保同时也安装了正确版本的 GPU 驱动。本文将介绍两种方法来检查 PyTorch 是否成功安装了 GPU 驱动。

方法一：使用 torch.cuda.is_available() 函数

PyTorch 中有一个函数 torch.cuda.is_available() 用于检查是否成功安装了 GPU 驱动。使用该函数的方法非常简单。代码如下所示：

import torch

if torch.cuda.is_available():
    print("Congratulations! You have successfully installed GPU version of PyTorch.")
else:
    print("Sorry, GPU version of PyTorch failed to install.")

运行上述代码后，程序将打印出一条消息，告诉你是否成功安装了 GPU 版本的 PyTorch。如果你看到了 “Congratulations!” 的消息，那么就代表你已经成功地安装了 GPU 驱动。

方法二：使用 nvidia-smi 命令

在 Linux 或者 Mac 上安装了 NVIDIA 显卡驱动之后，一般也会自动安装 NVIDIA System Management Interface (nvidia-smi) 。该命令可以显示当前已安装的 NVIDIA 显卡以及它们的显存使用情况。

要使用 nvidia-smi 命令，直接在终端输入该命令即可。如果该命令可以正常运行，并且输出显示了与你安装的显卡相匹配的设备信息和显存使用情况，那么就说明你的 GPU 驱动已经成功安装了。

以下是一个例子：

$ nvidia-smi
Sat Jul 31 13:26:43 2021       
+-----------------------------------------------------------------------------+
| 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. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  Off  | 00000000:03:00.0 Off |                    0 |
| N/A   34C    P0    42W / 300W |      0MiB / 16130MiB |      0%      Default |
|-------------------------------+----------------------+----------------------+
|   1  Tesla V100-SXM2...  Off  | 00000000:04:00.0 Off |                    0 |
| N/A   39C    P0    42W / 300W |      0MiB / 16130MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|

从上述输出可以看到，该计算机上有两个 NVIDIA Tesla V100-SXM2 显卡，在本例中它们均被关闭，显存使用率均为零。如果你也能看到类似这样的输出，那么就说明你的 GPU 驱动已经成功安装了。

总之，使用上述两种方法中的一个能够帮助你快速地检查 PyTorch 是否成功安装了 GPU 驱动。