1. 安装相关cuda cudnn
cuda是 11.6, 下载相应的离线包
sudo sh cuda.run
continue->accept->将驱动x按enter去掉,不按驱动了->不update软链接(可以多个cuda,不影响前面的)
同时复制相应的cudnn
在~/.bashrc中加入
export CUDA_HOME=/usr/local/cuda-11.6
export PATH=$PATH:$CUDA_HOME/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CUDA_HOME/lib64
2. pip安装相关pytorch安装教程
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
我试的用conda安装不好使
确定是否torch支持cuda,安装正确是true
print("gpu", torch.cuda.is_available())
3. extension相关错误
this扩展会有错误参考this缺少导致扩展pip错误