验证 Anaconda 和 Mamba 安装的 PyTorch 和 CUDA 是否正确

标签

AI

AI/硬件/GPU/NVIDIA

AI/硬件/GPU/NVIDIA/GPU

开发/Python/PyTorch

AI/算法/CUDA

开发/AI/CUDA

字数

143 字

阅读时间

1 分钟

参考安装清单：

channels:
  - pytorch
  - nvidia
  - conda-forge
dependencies:
  - pytorch
  - torchvision
  - torchaudio
  - pytorch-cuda=12.1
  - cuda-nvcc

检查 PyTorch

启动 python

python

>>> import torch
>>> torch.version.cuda
'12.1'
>>> torch.backends.cudnn.version()
8902
>>> torch.cuda.is_available()
True

或者撰写 Python 脚本和 Notebook：

import sys
import torch

print(sys.prefix)
print(torch.version.cuda)
print(torch.backends.cudnn.version())
print(torch.__version__)
print(torch.cuda.is_available())

检查 NVCC

运行 nvcc 命令检查是否正确安装 CUDA 编译器

nvcc --version

nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Fri_Dec_17_18:16:03_PST_2021
Cuda compilation tools, release 11.6, V11.6.55
Build cuda_11.6.r11.6/compiler.30794723_0

检查 PyTorch 拓展

>>> import torch.utils.cpp_extension
>>> torch.utils.cpp_extension.CUDA_HOME

贡献者

絢香猫

页面历史

最后编辑于 2 个月前

查看完整历史

4f5cdd7-docs: 添加文档