插件窝 干货文章 怎么看linux的pytorch环境

怎么看linux的pytorch环境

查看 torch CUDA print 710    来源:    2025-04-06

查看Linux下的PyTorch环境

要查看Linux系统中安装的PyTorch环境及相关配置,可以使用以下几种方法:

1. 使用Python交互式环境查看

import torch

# 查看PyTorch版本
print(torch.__version__)

# 查看CUDA是否可用
print(torch.cuda.is_available())

# 查看CUDA版本
print(torch.version.cuda)

# 查看cuDNN版本
print(torch.backends.cudnn.version())

# 查看当前设备
print(torch.cuda.current_device())

# 查看设备数量
print(torch.cuda.device_count())

# 查看当前设备名称
print(torch.cuda.get_device_name(0))

2. 使用conda或pip查看已安装包

# 使用conda查看
conda list | grep torch

# 使用pip查看
pip list | grep torch

3. 检查系统CUDA版本

# 查看系统CUDA版本
nvcc --version

# 或
cat /usr/local/cuda/version.txt

4. 检查cuDNN版本

# 查看cuDNN版本
cat /usr/local/cuda/include/cudnn.h | grep CUDNN_MAJOR -A 2

5. 查看GPU信息

# 查看NVIDIA显卡信息
nvidia-smi

# 查看所有GPU设备
lspci | grep -i nvidia

6. 创建简单的测试脚本

可以创建一个简单的Python脚本来测试PyTorch是否能正常工作:

import torch

x = torch.rand(5, 3)
print(x)

if torch.cuda.is_available():
    print("CUDA is available!")
    device = torch.device("cuda")
    y = torch.ones_like(x, device=device)
    x = x.to(device)
    z = x + y
    print(z)
    print(z.to("cpu", torch.double))
else:
    print("CUDA is NOT available!")

注意事项

  1. 确保你是在正确的Python环境中运行这些命令(特别是使用conda或virtualenv时)
  2. 如果使用Docker容器,需要在容器内执行这些命令
  3. 不同版本的PyTorch可能有略微不同的API

通过这些方法,你可以全面了解Linux系统中PyTorch环境的配置情况,包括版本信息、CUDA支持情况以及GPU可用性等。