共计 3657 个字符,预计需要花费 10 分钟才能阅读完成。
环境准备
在开始之前,确保你的系统满足以下要求:

- Ubuntu 22.04 LTS(推荐使用最新稳定版)
- 至少 16GB 内存(32GB 以上更佳)
- 足够的磁盘空间(建议 100GB 以上)
- 稳定的互联网连接
驱动版本选择
RTX 5090 需要 NVIDIA 官方最新驱动支持。目前推荐使用 525 系列或更高版本驱动。可以通过以下命令查看可用驱动版本:
ubuntu-drivers devices
详细安装步骤
1. NVIDIA 驱动安装
首先禁用 nouveau 驱动,这是 Linux 默认的开源 NVIDIA 驱动,可能与官方驱动冲突:
sudo bash -c "echo'blacklist nouveau'>> /etc/modprobe.d/blacklist-nvidia-nouveau.conf"
sudo bash -c "echo'options nouveau modeset=0'>> /etc/modprobe.d/blacklist-nvidia-nouveau.conf"
sudo update-initramfs -u
然后重启系统。重启后,添加官方 PPA 源并安装驱动:
sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
sudo apt install nvidia-driver-525
安装完成后,重启系统并验证驱动是否安装成功:
nvidia-smi
2. CUDA 12.x 工具链安装
CUDA 是 NVIDIA 提供的并行计算平台和编程模型。安装 CUDA 12.x:
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda-repo-ubuntu2204-12-1-local_12.1.1-530.30.02-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-12-1-local_12.1.1-530.30.02-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-12-1-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda
安装完成后,将 CUDA 添加到环境变量:
export PATH=/usr/local/cuda-12.1/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
3. cuDNN 配置与验证
cuDNN 是 NVIDIA 提供的深度神经网络库。首先从 NVIDIA 开发者网站下载对应版本的 cuDNN,然后安装:
sudo dpkg -i libcudnn8_8.9.0.131-1+cuda12.1_amd64.deb
sudo dpkg -i libcudnn8-dev_8.9.0.131-1+cuda12.1_amd64.deb
sudo dpkg -i libcudnn8-samples_8.9.0.131-1+cuda12.1_amd64.deb
验证 cuDNN 安装:
cp -r /usr/src/cudnn_samples_v8/ $HOME
cd $HOME/cudnn_samples_v8/mnistCUDNN
make clean && make
./mnistCUDNN
如果看到 ”Test passed!”,说明 cuDNN 安装成功。
深度学习框架安装
PyTorch with CUDA 支持
推荐使用 conda 创建虚拟环境后安装 PyTorch:
conda create -n pytorch_env python=3.9
conda activate pytorch_env
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
TensorFlow GPU 版本
在同一个或另一个 conda 环境中安装 TensorFlow GPU 版本:
conda create -n tf_env python=3.9
conda activate tf_env
conda install -c conda-forge cudatoolkit=12.1.1
pip install tensorflow-gpu
验证与测试
PyTorch 测试脚本
import torch
# 检查 CUDA 是否可用
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"CUDA device count: {torch.cuda.device_count()}")
print(f"Current device: {torch.cuda.current_device()}")
print(f"Device name: {torch.cuda.get_device_name(0)}")
# 简单的矩阵运算测试
a = torch.randn(10000, 10000).cuda()
b = torch.randn(10000, 10000).cuda()
c = torch.matmul(a, b)
print(f"Matrix multiplication completed: {c.shape}")
TensorFlow 测试脚本
import tensorflow as tf
# 检查 GPU 是否可用
print(f"GPU available: {tf.config.list_physical_devices('GPU')}")
# 简单的矩阵运算测试
with tf.device('/GPU:0'):
a = tf.random.normal((10000, 10000))
b = tf.random.normal((10000, 10000))
c = tf.matmul(a, b)
print(f"Matrix multiplication completed: {c.shape}")
性能基准测试
可以使用以下命令测试 GPU 性能:
# 安装基准测试工具
pip install gpustat
# 实时监控 GPU 使用情况
gpustat -i
# 运行官方基准测试
cd /usr/local/cuda/samples/1_Utilities/bandwidthTest
make
./bandwidthTest
常见问题解决
驱动冲突处理
如果遇到驱动冲突,可以尝试以下步骤:
sudo apt purge nvidia*
sudo apt autoremove
sudo ubuntu-drivers autoinstall
CUDA 版本兼容性问题
如果遇到 CUDA 版本不兼容,可以安装多个 CUDA 版本并使用 update-alternatives 管理:
sudo update-alternatives --install /usr/local/cuda cuda /usr/local/cuda-12.1 121
sudo update-alternatives --config cuda
多 GPU 配置注意事项
对于多 GPU 系统,确保 PCIe 带宽足够,并在代码中正确设置设备:
# PyTorch 多 GPU 示例
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 或者使用 DataParallel
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
最佳实践
环境隔离方案
推荐使用 conda 或 docker 进行环境隔离:
# 使用 conda
conda create -n dl_env python=3.9
conda activate dl_env
# 使用 docker
sudo docker run --gpus all -it nvidia/cuda:12.1.1-base-ubuntu22.04
性能优化建议
- 使用混合精度训练
- 优化数据加载流程
- 合理设置 batch size
- 使用 CUDA Graphs 减少内核启动开销
- 定期更新驱动和 CUDA 版本
总结
本文详细介绍了在 Ubuntu 22.04 系统上为 RTX 5090 配置深度学习环境的完整流程。从驱动安装到框架配置,再到性能测试和优化,涵盖了实际使用中的所有关键步骤。按照本指南操作后,你应该能够充分利用 RTX 5090 的强大计算能力进行深度学习研究和开发。
