Ubuntu下配置PyTorch

一、下载cuda和cudnn

二、配置conda

三、下载pytorch

一、下载cuda和cudnn

1、打开命令提示符输入 nvidia-smi 查看版本

2、cuda：CUDA Toolkit Archive | NVIDIA Developer

根据自己的需求选择对应的版本

根据给出命令依次执行，最后输入 nvcc -V ，出现版本号则安装成功。

添加环境变量

复制代码

vi ~/.bashrc

#复制到文件末尾，记得替换自己安装路径，可以通过运行whereis cuda查看cuda路径

export PATH=/usr/local/cuda-12.6/bin${PATH:+:${PATH}}

export LD_LIBRARY_PATH=/usr/local/cuda-12.6/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

重新加载环境配置文件

复制代码

source ~/.bashrc

3、cudnn：cuDNN Archive | NVIDIA Developer

我是wsl下的ubuntu，所以直接将文件复制到/home/user/下，然后进行解压

复制代码

tar -xvf cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz

然后执行下述命令，记得修改成自己的路径

复制代码

sudo cp cudnn-linux-x86_64-8.9.7.29_cuda12-archive/include/* /usr/local/cuda/include

sudo cp cudnn-linux-x86_64-8.9.7.29_cuda12-archive/lib/libcudnn* /usr/local/cuda/lib64

sudo chmod a+r /usr/local/cuda/include/cudnn.h /usr/local/cuda/lib64/libcudnn*

查看cudnn的版本

复制代码

cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

出现下图内容则安装成功

二、配置conda

复制代码

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh

bash Miniconda3-latest-Linux-x86_64.sh -b -p $HOME/miniconda3

source ~/miniconda3/bin/activate

conda init bash

验证

复制代码

conda --version # ≥24.x

conda create -n py310 python=3.10  #若出现选择，则全部选accept(a)，输入a即可

若验证成功，根据提示激活虚拟环境

复制代码

conda activate py310

三、下载pytorch

在虚拟环境py310中执行以下命令

复制代码

pip install --no-cache-dir torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126

#下载慢可选择镜像
pip install torch==2.6.0+cu126 torchvision==0.21.0+cu126 torchaudio==2.6.0+cu126 -f https://mirror.sjtu.edu.cn/pytorch-wheels/torch_stable.html

验证

复制代码

python -c "import torch, sys, platform;\

print('PyTorch', torch.__version__);\

print('CUDA ', torch.version.cuda);\

print('GPU ', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None')"

参考：

https://zhuanlan.zhihu.com/p/691711768

https://www.cnblogs.com/HorizonTree/articles/18806799