Ubuntu PyTorch如何配置GPU - 问答

在Ubuntu上配置PyTorch以使用GPU，你需要确保你的系统满足以下条件：

NVIDIA GPU：首先，你需要一个支持CUDA的NVIDIA GPU。
NVIDIA驱动：安装适用于你的GPU型号的最新NVIDIA驱动。
CUDA Toolkit：安装与你的GPU和PyTorch版本兼容的CUDA Toolkit。
cuDNN库：安装与你的CUDA版本兼容的cuDNN库。
NCCL库（可选）：如果你打算进行分布式训练，你可能还需要安装NCCL库。
PyTorch：安装支持GPU的PyTorch版本。

以下是配置步骤：

步骤 1: 安装NVIDIA驱动

首先，你需要安装NVIDIA驱动。你可以从NVIDIA官网下载适合你GPU型号的驱动程序，并按照说明进行安装。

步骤 2: 安装CUDA Toolkit

访问NVIDIA的CUDA Toolkit下载页面，选择适合你的Ubuntu版本和GPU架构的CUDA Toolkit版本进行安装。你可以使用以下命令来安装CUDA Toolkit：

# 添加NVIDIA CUDA PPA
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

# 更新包列表并安装CUDA Toolkit
sudo apt-get update
sudo apt-get install -y cuda

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

步骤 3: 安装cuDNN库

你需要注册一个NVIDIA开发者账户来下载cuDNN。下载后，解压缩文件并将文件复制到CUDA目录中：

# 假设你已经下载了cuDNN压缩包到~/Downloads目录
tar -xzvf cudnn-*-linux-x64-v*.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: 安装NCCL库（可选）

如果你需要进行分布式训练，你需要安装NCCL库。你可以从NVIDIA官网下载NCCL并按照说明进行安装。

步骤 5: 安装PyTorch

你可以使用pip或conda来安装PyTorch。确保选择与你的CUDA版本兼容的PyTorch版本。例如，使用pip安装支持CUDA 11.3的PyTorch：

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

或者，如果你使用conda，可以使用以下命令：

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

步骤 6: 验证安装

安装完成后，你可以通过运行以下命令来验证PyTorch是否能检测到GPU：

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

如果输出显示了你的GPU型号并且torch.cuda.is_available()返回True，那么恭喜你，PyTorch已经成功配置为使用GPU了！

0 赞

0 踩