原理

详细原理参见https://www.cnblogs.com/yhjoker/p/10972795.html

简单原理就是通过修改$PATH环境变量，pytorch会以最前面的CUDA版本为准

例如，你可能装过2个版本的cuda

/usr/local/cuda-11.1
/usr/local/cuda-11.0


查看以下两个环境变量

echo $PATH
echo $LD_LIBRARY_PATH


可以看到我的此shell下，cuda11.0在cuda11.1前面，那么pytorch实际运行的就是11.0版本的cuda

修改方法

永久修改当前用户下的环境变量

1. 修改 ~/.bashrc 文件

vim ~/.bashrc


2. 添加

export LD_LIBRARY_PATH=/usr/local/cuda-11.0/lib64:$LD_LIBRARY_PATH
export PATH=/usr/local/cuda-11.0/bin:$PATH


3. 在当前shell下重新运行 ~/.bashrc 文件

source ~/.bashrc


这样每次打开新shell的时候会自动运行~/.bashrc文件，以达到永久性修改。（不过最新的chrome ssh安全外壳有毒，需要手动运行一下。。）

不用担心来回修改版本，source了太多次~/.bashrc文件添加了很多环境变量，你新开一个shell就是了

临时修改

不多说了，直接运行加的那两行代码即可

export LD_LIBRARY_PATH=/usr/local/cuda-11.0/lib64:$LD_LIBRARY_PATH
export PATH=/usr/local/cuda-11.0/bin:$PATH


检验效果

老方法，使用pytorch的方法查看

torch.version.cuda


你突然发现，纳尼，怎么还是老版本cuda没变?



事实上，上述输出的 cuda 的版本并不一定是 Pytorch 在实际系统上运行时使用的 cuda 版本，而是编译该 Pytorch release 版本时使用的 cuda 版本。

要查看运行时使用的 cuda 版本，输入以下代码

import torch.utils.cpp_extension
torch.utils.cpp_extension.CUDA_HOME


发现 Pytorch 运行时使用的 cuda 改了