Ubuntu服务器安装CUDA和PyTorch的步骤是什么?

以下是 Ubuntu 服务器安装 CUDA 和 PyTorch 的详细步骤:

一、准备工作

1. 检查系统要求

# 查看系统信息
lsb_release -a
uname -m

# 检查GPU信息
lspci | grep -i nvidia

2. 安装依赖

sudo apt update
sudo apt upgrade -y
sudo apt install build-essential gcc make perl dkms -y

二、安装NVIDIA驱动

方法1:使用官方驱动(推荐)

# 添加显卡驱动PPA
sudo add-apt-repository ppa:graphics-drivers/ppa -y
sudo apt update

# 查看推荐驱动版本
ubuntu-drivers devices

# 安装推荐驱动(或指定版本)
sudo apt install nvidia-driver-550 -y  # 根据推荐选择版本

# 重启系统
sudo reboot

方法2:使用CUDA Toolkit包含的驱动

# 如果使用CUDA Toolkit安装,它会自动安装对应驱动

三、安装CUDA Toolkit

方法1:使用官方runfile(推荐)

  1. 访问 NVIDIA CUDA下载页面
  2. 选择对应版本,获取安装命令,例如:
    wget https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run
    sudo sh cuda_12.4.0_550.54.14_linux.run

方法2:使用网络安装

# 下载网络安装包
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update
sudo apt install cuda-toolkit-12-4 -y  # 选择需要的版本

四、配置环境变量

# 编辑bashrc
nano ~/.bashrc

# 添加以下内容
export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
export CUDA_HOME=/usr/local/cuda

# 使配置生效
source ~/.bashrc

五、验证CUDA安装

# 检查驱动版本
nvidia-smi

# 检查CUDA版本
nvcc --version

# 运行CUDA示例测试
cd /usr/local/cuda/samples/1_Utilities/deviceQuery
sudo make
./deviceQuery

六、安装PyTorch

方法1:使用pip安装(推荐)

# 创建虚拟环境(可选但推荐)
sudo apt install python3-venv -y
python3 -m venv pytorch_env
source pytorch_env/bin/activate

# 安装PyTorch(访问 https://pytorch.org 获取最新命令)
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

方法2:使用conda安装

# 安装Miniconda
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

# 创建环境并安装PyTorch
conda create -n pytorch_env python=3.10
conda activate pytorch_env
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

七、验证PyTorch安装

import torch

# 检查PyTorch版本
print(torch.__version__)

# 检查CUDA是否可用
print(torch.cuda.is_available())

# 检查CUDA版本
print(torch.version.cuda)

# 检查GPU数量
print(torch.cuda.device_count())

# 检查当前GPU
print(torch.cuda.current_device())

# 检查GPU名称
print(torch.cuda.get_device_name(0))

# 简单张量计算测试
x = torch.rand(5, 3).cuda()
print(x)

八、安装cuDNN(可选但推荐)

# 从NVIDIA开发者网站下载cuDNN
# 需要注册NVIDIA开发者账号
# 下载对应CUDA版本的cuDNN

# 安装示例(具体文件名根据下载版本调整)
sudo dpkg -i libcudnn8_8.9.7.29-1+cuda12.2_amd64.deb
sudo dpkg -i libcudnn8-dev_8.9.7.29-1+cuda12.2_amd64.deb

九、常见问题解决

1. 驱动安装失败

# 卸载原有驱动
sudo apt purge nvidia-*
sudo apt autoremove

# 禁用nouveau驱动
sudo nano /etc/modprobe.d/blacklist-nouveau.conf
# 添加:blacklist nouveau
# 添加:options nouveau modeset=0

sudo update-initramfs -u
sudo reboot

2. CUDA版本不匹配

# 查看PyTorch支持的CUDA版本
# 访问 https://pytorch.org 查看兼容性矩阵

# 安装指定版本的PyTorch
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121

十、一键安装脚本示例

#!/bin/bash
# 保存为 install_cuda_pytorch.sh

echo "正在更新系统..."
sudo apt update && sudo apt upgrade -y

echo "正在安装依赖..."
sudo apt install build-essential gcc make perl dkms python3-pip -y

echo "正在安装NVIDIA驱动..."
sudo add-apt-repository ppa:graphics-drivers/ppa -y
sudo apt update
sudo apt install nvidia-driver-550 -y

echo "请重启系统后继续执行第二部分脚本"

注意事项:

  1. 确保CUDA版本与PyTorch版本兼容
  2. 建议使用虚拟环境管理Python包
  3. 服务器安装时注意网络连接稳定性
  4. 多GPU环境需要额外配置NCCL
  5. 生产环境建议使用Docker容器化部署

根据你的具体需求(如CUDA版本、PyTorch版本),可能需要调整上述命令中的版本号。

云服务器