当前位置：首页 > news >正文

【笔记】NVIDIA AI Workbench 中安装 PyTorch

news 来源：原创 2025/6/13 9:47:48

一、前置条件

（部分整理，未完整罗列）

1、WSL-NVIDIA-Workbench 中安装 Anaconda
- 确保已经安装了 Anaconda 并配置好 Python 环境。
2、NVIDIA AI Workbench 中 sudo 密码问题排查与解决
- 解决任何可能的权限问题，确保可以顺利执行需要 root 权限的操作。
3、NVIDIA AI Workbench 中安装 CUDA 12.9
- 确保 CUDA 已正确安装，并且版本符合要求。CUDA 是 PyTorch 使用 GPU 加速的基础。
4、NVIDIA AI Workbench 中安装 cuDNN 9.10.2
- 确保 cuDNN 已正确安装，并且版本与 CUDA 兼容。cuDNN 提供了优化的深度神经网络算法。
5、确保网络通畅装

【深度学习环境搭建】WSL-NVIDIA-Workbench 中安装 Anaconda-CSDN博客

【笔记】NVIDIA AI Workbench 中 sudo 密码问题排查与解决-CSDN博客

【笔记】NVIDIA AI Workbench 中安装 CUDA 12.9-CSDN博客

【笔记】NVIDIA AI Workbench 中安装 cuDNN 9.10.2-CSDN博客

二、安装 PyTorch

（一）打开官网

访问 PyTorch 官网，找到“快速开始”或“安装”部分。

（二）选择安装命令

由于我的系统已安装 CUDA 12.9（以 cu128举例），你可以使用以下命令安装 PyTorch：

预发布版

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128

这条命令会安装预发布版的 PyTorch，支持 CUDA 12.8。如果我们希望安装稳定版，可以切换官对应的稳定版安装命令。

稳定版

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

三、验证 PyTorch 安装

安装完成后，可以通过以下 Python 脚本验证 PyTorch 是否正确安装并能够利用 CUDA 和 cuDNN：

先输入：

python

再整段复制粘贴，回车运行：

import torch  # 导入 PyTorch 库print("PyTorch 版本：", torch.__version__)  # 打印 PyTorch 的版本号# 检查 CUDA 是否可用，并设置设备（"cuda:0" 或 "cpu"）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("设备：", device)  # 打印当前使用的设备
print("CUDA 可用：", torch.cuda.is_available())  # 打印 CUDA 是否可用
print("cuDNN 已启用：", torch.backends.cudnn.enabled)  # 打印 cuDNN 是否已启用# 打印 PyTorch 支持的 CUDA 和 cuDNN 版本
print("支持的 CUDA 版本：", torch.version.cuda)
print("cuDNN 版本：", torch.backends.cudnn.version())# 创建两个随机张量（默认在 CPU 上）
x = torch.rand(5, 3)
y = torch.rand(5, 3)# 将张量移动到指定设备（CPU 或 GPU）
x = x.to(device)
y = y.to(device)# 对张量进行逐元素相加
z = x + y# 打印结果
print("张量 z 的值：")
print(z)  # 输出张量 z 的内容

四、预期输出

验证PyTorch深度学习环境Torch和CUDA还有cuDNN是否正确配置的命令-CSDN博客

如果一切正常，我们应该看到类似以下的输出：

>>> print("PyTorch 版本：", torch.__version__)  # 打印 PyTorch 的版本号
PyTorch 版本： 2.8.0.dev20250610+cu128
>>>
>>> # 检查 CUDA 是否可用，并设置设备（"cuda:0" 或 "cpu"）
>>> device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
>>> print("设备：", device)  # 打印当前使用的设备
设备： cuda:0
>>> print("CUDA 可用：", torch.cuda.is_available())  # 打印 CUDA 是否可用
CUDA 可用： True
>>> print("cuDNN 已启用：", torch.backends.cudnn.enabled)  # 打印 cuDNN 是否已启用
cuDNN 已启用： True
>>>
>>> # 打印 PyTorch 支持的 CUDA 和 cuDNN 版本
>>> print("支持的 CUDA 版本：", torch.version.cuda)
支持的 CUDA 版本： 12.8
>>> print("cuDNN 版本：", torch.backends.cudnn.version())
cuDNN 版本： 90800
>>>
>>> # 创建两个随机张量（默认在 CPU 上）
>>> x = torch.rand(5, 3)
>>> y = torch.rand(5, 3)
>>>
>>> # 将张量移动到指定设备（CPU 或 GPU）
>>> x = x.to(device)
>>> y = y.to(device)
>>>
>>> # 对张量进行逐元素相加
>>> z = x + y
>>>
>>> # 打印结果
>>> print("张量 z 的值：")
张量 z 的值：
>>> print(z)  # 输出张量 z 的内容
tensor([[1.2157, 1.3211, 0.8185],[0.9087, 1.0803, 0.6371],[1.5480, 1.5786, 1.0636],[1.6989, 1.2323, 1.5848],[0.9905, 1.6324, 1.5171]], device='cuda:0')
>>>

这表明 PyTorch 已经成功安装，并且能够利用 CUDA 和 cuDNN 进行 GPU 计算。