检查 CUDA 版本 PyTorch - Python

如果你正在使用 PyTorch 通过 CUDA 运行你的深度学习模型，确保你正确配置了 CUDA 和 PyTorch 版本非常重要。本文将介绍如何通过 Python 代码检查你的 CUDA 版本以及 PyTorch 是否正确配置了 CUDA。

检查 CUDA 版本

要检查你的 CUDA 版本，可以运行以下代码：

import torch

print(torch.version.cuda)

这将打印出你正在使用的 CUDA 版本。

检查 PyTorch 是否正确配置了 CUDA

要检查 PyTorch 是否正确配置了 CUDA，可以运行以下代码：

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

print('Using device:', device)
print('CUDA version:', torch.version.cuda)
print('PyTorch version:', torch.__version__)

x = torch.rand(5, 3).to(device)
y = torch.rand(5, 3).to(device)

print(x + y)

这将打印出你的设备类型、CUDA 版本和 PyTorch 版本，并在 CUDA 设备上随机初始化两个矩阵并将它们相加。

如果一切配置正确，你应该看到如下输出：

Using device: cuda
CUDA version: 10.1
PyTorch version: 1.9.0
tensor([[1.1932, 1.3653, 0.7888],
        [0.6227, 0.9866, 0.4831],
        [0.6355, 1.1285, 0.6019],
        [1.1166, 1.2428, 0.8734],
        [0.7825, 1.6829, 0.6225]], device='cuda:0')

如果 CUDA 未正确配置，你将收到一条错误消息。如果你的 PyTorch 版本未与你的 CUDA 版本兼容，你也会看到类似的错误消息。在这种情况下，你应该更新 PyTorch 到与所安装的 CUDA 版本兼容的最新版本。

结论

在运行深度学习模型之前，始终检查你的 CUDA 版本和 PyTorch 是否正确配置了 CUDA 非常重要。运行上面的代码片段，你将能够立即检查你的配置。