在 PyTorch 中,你可以很容易地检查 CUDA 是否可用、CUDA 的版本以及 cuDNN 的版本信息。以下是如何在 PyTorch 中获取这些信息的步骤:
检查 CUDA 是否可用
首先,你需要检查 CUDA 是否在你的系统上可用。PyTorch 提供了一个简单的方法来检查这一点。
python
import torch
# 检查 CUDA 是否可用
print(f"CUDA available: {torch.cuda.is_available()}")
# 如果 CUDA 可用,打印当前设备名称
if torch.cuda.is_available():
print(f"Current device: {torch.cuda.get_device_name(0)}")
获取 CUDA 版本
你可以使用 torch.version.cuda
来获取 CUDA 的版本信息。
python
import torch
# 获取 CUDA 版本
cuda_version = torch.version.cuda
print(f"CUDA Version: {cuda_version}")
获取 cuDNN 版本
你可以使用 torch.backends.cudnn.version()
来获取 cuDNN 的版本信息。
python
import torch
# 获取 cuDNN 版本
cudnn_version = torch.backends.cudnn.version()
print(f"cuDNN Version: {cudnn_version}")
完整示例代码
将上述所有部分组合在一起,完整的示例代码如下:
python
import torch
# 检查 CUDA 是否可用
print(f"CUDA available: {torch.cuda.is_available()}")
# 如果 CUDA 可用,打印当前设备名称
if torch.cuda.is_available():
print(f"Current device: {torch.cuda.get_device_name(0)}")
# 获取 CUDA 版本
cuda_version = torch.version.cuda
print(f"CUDA Version: {cuda_version}")
# 获取 cuDNN 版本
cudnn_version = torch.backends.cudnn.version()
print(f"cuDNN Version: {cudnn_version}")
运行结果示例
假设你运行了上述代码,并且 CUDA 和 cuDNN 都已正确安装和配置,你可能会看到类似以下的输出:
plaintext
CUDA available: True
Current device: NVIDIA GeForce RTX 3090
CUDA Version: 11.1
cuDNN Version: 8200
注意事项
- 环境配置 :确保你的环境中已经安装了 PyTorch 并且支持 CUDA。如果你使用的是 CPU 版本的 PyTorch,
torch.cuda.is_available()
将返回False
。 - 依赖库:确保你已经安装了正确的 CUDA 和 cuDNN 库,并且它们与你的 PyTorch 版本兼容。
- 多 GPU 环境 :如果你有多个 GPU,可以使用
torch.cuda.device_count()
来获取 GPU 数量,并通过torch.cuda.get_device_name(i)
获取每个 GPU 的名称。
希望这些信息能帮助你检查 CUDA 和 cuDNN 的版本信息!如果你有更多问题或需要进一步的帮助,请告诉我。