Ubuntu上PyTorch与CUDA如何协同工作 - 问答

在Ubuntu系统上，PyTorch与CUDA的协同工作主要依赖于NVIDIA提供的CUDA工具包和驱动程序，以及PyTorch的CUDA支持库。以下是它们协同工作的基本步骤：

1. 安装NVIDIA驱动程序

首先，确保你的系统已经安装了兼容的NVIDIA显卡驱动程序。你可以从NVIDIA官方网站下载并安装最新的驱动程序。

sudo apt update
sudo apt install nvidia-driver-<version>

2. 安装CUDA工具包

接下来，安装CUDA工具包。你可以从NVIDIA官方网站下载适用于Ubuntu的CUDA Toolkit安装包，并按照官方指南进行安装。

# 下载CUDA Toolkit安装包
wget https://developer.download.nvidia.com/compute/cuda/<version>/cuda_<version>_linux.run

# 安装CUDA Toolkit
sudo sh cuda_<version>_linux.run

安装过程中，确保选择安装CUDA Toolkit和驱动程序。

3. 安装cuDNN库

cuDNN（CUDA Deep Neural Network library）是一个用于深度神经网络的GPU加速库。你需要从NVIDIA官方网站下载并安装cuDNN库。

# 下载cuDNN库
wget https://developer.nvidia.com/compute/machine-learning/cudnn/secure/<version>/prod/<version>_linux_x64_v<version>.tgz

# 解压并安装cuDNN库
tar -xzvf <version>_linux_x64_v<version>.tgz
sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

4. 安装PyTorch与CUDA支持

使用pip安装PyTorch，并确保安装的版本支持CUDA。

# 安装PyTorch与CUDA支持
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu<version>

将<version>替换为你安装的CUDA版本号。例如，如果你安装的是CUDA 11.7，命令应为：

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117

5. 验证安装

验证PyTorch是否能够检测到CUDA设备。

import torch

# 检查是否有可用的GPU
print(torch.cuda.is_available())

# 获取GPU设备数量
print(torch.cuda.device_count())

# 获取当前GPU设备名称
print(torch.cuda.get_device_name(0))

如果输出显示True，并且能够列出GPU设备名称，则说明PyTorch与CUDA已经成功协同工作。

6. 配置环境变量

确保CUDA和cuDNN的路径已经添加到系统的环境变量中。编辑~/.bashrc文件，添加以下内容：

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

然后运行以下命令使更改生效：

source ~/.bashrc

通过以上步骤，你可以在Ubuntu系统上成功安装并配置PyTorch与CUDA，使其能够协同工作。

0 赞

0 踩