在Linux环境下,PyTorch与CUDA的协同工作主要涉及以下几个步骤:
首先,你需要安装CUDA Toolkit,这是NVIDIA提供的用于深度学习的并行计算平台和API。你可以从NVIDIA官网下载适合你GPU型号的CUDA Toolkit。
# 下载CUDA Toolkit安装包(例如cuda_11.7.0_515.43.04_linux.run)
wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_515.43.04_linux.run
# 安装CUDA Toolkit
sudo sh cuda_11.7.0_515.43.04_linux.run
安装过程中,选择安装CUDA Toolkit和驱动程序。
cuDNN是NVIDIA提供的深度神经网络库,用于加速深度学习框架的计算。你需要下载并安装与CUDA Toolkit版本兼容的cuDNN库。
# 下载cuDNN(例如cudnn-11.7-linux-x64-v8.4.1.50.tgz)
wget https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.4.1.50/11.7_20210301/cudnn-11.7-linux-x64-v8.4.1.50.tgz
# 解压cuDNN
tar -xzvf cudnn-11.7-linux-x64-v8.4.1.50.tgz
# 复制cuDNN文件到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*
为了使系统能够找到CUDA和cuDNN库,你需要配置环境变量。
# 编辑~/.bashrc文件
echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
# 使环境变量生效
source ~/.bashrc
你可以使用pip或conda来安装支持CUDA的PyTorch版本。
# 安装支持CUDA 11.7的PyTorch
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
# 创建一个新的conda环境(可选)
conda create -n pytorch_env python=3.9
# 激活环境
conda activate pytorch_env
# 安装支持CUDA 11.7的PyTorch
conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch
安装完成后,你可以验证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协同工作。
通过以上步骤,你可以在Linux环境下成功安装并配置PyTorch与CUDA,从而利用GPU加速深度学习模型的训练和推理。