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

验证

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