Ubuntu下PyTorch与CUDA如何配合

在Ubuntu系统下，PyTorch与CUDA的配合使用可以显著提升深度学习模型的训练和推理速度。以下是详细的步骤指南：

1. 安装CUDA Toolkit

首先，你需要安装CUDA Toolkit。你可以从NVIDIA官方网站下载适合你显卡的CUDA Toolkit版本。

下载CUDA Toolkit

访问 NVIDIA CUDA Toolkit下载页面。
选择适合你Ubuntu版本的CUDA Toolkit版本。
下载并运行安装脚本。

安装步骤

# 添加NVIDIA PPA
sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt-get update

# 安装NVIDIA驱动（如果尚未安装）
sudo apt-get install nvidia-driver-<版本号>

# 安装CUDA Toolkit
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget http://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2004-11-8-local_11.8.0-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2004-11-8-local_11.8.0-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-8-local/7fa2af80.pub
sudo apt-get update
sudo apt-get -y install cuda

# 设置环境变量
echo 'export PATH=/usr/local/cuda-11.8/bin${PATH:+:${PATH}}' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
source ~/.bashrc

2. 安装cuDNN

cuDNN是NVIDIA提供的深度神经网络库，可以进一步加速深度学习模型的训练。

下载cuDNN

访问 NVIDIA cuDNN下载页面。
注册并登录NVIDIA开发者账号。
下载适合你CUDA版本的cuDNN库。

安装步骤

# 下载cuDNN库
tar -xzvf cudnn-11.8-linux-x64-v8.6.0.26.tgz

# 复制文件到CUDA目录
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*

3. 安装PyTorch

你可以使用pip来安装PyTorch，并指定CUDA版本。

安装步骤

# 创建一个新的虚拟环境（可选）
python3 -m venv pytorch_env
source pytorch_env/bin/activate

# 安装PyTorch with CUDA支持
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118

4. 验证安装

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

import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.current_device())
print(torch.cuda.get_device_name(torch.cuda.current_device()))

如果输出显示True和你的GPU型号，说明PyTorch已经成功配置并可以使用CUDA。

5. 使用PyTorch与CUDA

在编写PyTorch代码时，你可以直接使用GPU进行计算。

import torch

# 创建一个张量并将其移动到GPU
x = torch.tensor([1.0, 2.0, 3.0], device='cuda')

# 创建另一个张量并进行计算
y = x + 1

print(y)

通过以上步骤，你可以在Ubuntu系统下成功配置和使用PyTorch与CUDA，从而加速深度学习模型的训练和推理。