在使用 Pytorch 时,由于 Pytorch 和 cuda 版本的更新,可能出现程序运行时需要特定版本的 cuda 进行运行环境支持的情况,如使用特定版本的 cuda 编译 CUDAExtension 引入的拓展模块等。为了满足应用程序和框架本身对不同版本的 cuda 的需求,Pytorch 需要能够在不同版本的 cuda 之间切换使用。这里即记录 Pytorch 在编辑 cpp 和 cuda 拓展时确定所使用 cuda 版本的基本流程以及 Pytorch 使用不同版本的 cuda 运行程序的方法。
Pytorch 确定所使用的 cuda 版本
实际使用过程中,Pytorch 检测运行时使用的 cuda 版本的代码位于 torch/utils/cpp_extension.py 的_find_cuda_home 函数 ( Pytorch 1.1.0, )中.这里主要介绍 Linux 环境下的 cuda 版本的确认过程,关于 Windows 环境下多版本 cuda 的使用可以参考上述文件中的具体实现.
确定 cuda 路径
若在运行时需要使用 cuda 进行程序的编译或其他 cuda 相关的操作,Pytorch 会首先定位一个 cuda 安装目录( 来获取所需的特定版本 cuda 提供的可执行程序、库文件和头文件等文件 )。具体而言,Pytorch 首先尝试获取环境变量 CUDA_HOME/CUDA_PATH 的值作为运行时使用的 cuda 目录。若直接设置了 CUDA_HOME/CUDA_PATH 变量,则 Pytorch 使用 CUDA_HOME/CUDA_PATH 指定的路径作为运行时使用的 cuda 版本的目录。
若上述环境变量不存在,则 Pytorch 会检查系统是否存在固定路径 /usr/local/cuda 。默认情况下,系统并不存在对环境变量 CUDA_HOME 设置,故而 Pytorch 运行时默认检查的是 Linux 环境中固定路径 /usr/local/cuda 所指向的 cuda 目录。 /usr/local/cuda 实际上是一个软连接文件,当其存在时一般被设置为指向系统中某一个版本的 cuda 文件夹。使用一个固定路径的软链接的好处在于,当系统中存在多个安装的 cuda 版本时,只需要修改上述软连接实际指向的 cuda 目录,而不需要修改任何其他的路径接口,即可方便的通过唯一的路径使用不同版本的 cuda. 如笔者使用的服务器中,上述固定的 /usr/local/cuda 路径即指向一个较老的 cuda-8.0 版本的目录。
需要注意的是, /usr/local/cuda 并不是一个 Linux 系统上默认存在的路径,其一般在安装 cuda 时创建( 为可选项,不强制创建 )。故而 Pytorch 检测上述路径时也可能会失败。
若 CUDA_HOME 变量指定的路径和默认路径 /usr/local/cuda 均不存在安装好的 cuda 目录,则 Pytorch 通过运行命令 which nvcc 来找到一个包含有 nvcc 命令的 cuda 安装目录,并将其作为运行时使用的 cuda 版本。具体而言,系统会根据环境变量 PATH 中的目录去依次搜索可用的 nvcc 可执行文件,若环境变量 PATH 中包含多个安装好的 cuda 版本的可执行文件目录( 形如/home/test/cuda-10.1/bin ),则排在 PATH 中的第一个 cuda 的可执行文件目录中的 nvcc 命令会被选中,其所对应的路径被选为 Pytorch 使用的 cuda 路径。同样的,若 PATH 中不存在安装好的 cuda 版本的可执行目录,则上述过程会失败,Pytorch 最终会由于找不到可用的 cuda 目录而无法使用 cuda.比较推荐的做法是保持 PATH 路径中存在唯一一个对应所需使用的 cuda 版本的可执行目录的路径。
在确定好使用的 cuda 路径后,基于 cuda 的 Pytorch 拓展即会使用确定好的 cuda 目录中的可执行文件( /bin )、头文件( /include )和库文件( /lib64 )完成所需的编译过程。
sudo rm -rf /usr/local/cuda //删除软链接,注意是 /usr/local/cuda 而不是 /usr/local/cuda/,前者仅删除软链接,而后者会删除软链接所指向的目录的所有内容,操作请小心 sudo ln -s cuda_path /usr/local/cuda //创建名为 /usr/local/cuda 的软链接,其指向 cuda_path 所指定的 cuda 安装目录
sudo ln -sf cuda_path /usr/local/cuda //修改或创建软链接 /usr/local/cuda 使其指向指定版本的 cuda 目录
export CUDA_HOME=/home/test/cuda-10.1/ //设置全局变量 CUDA_HOME export PATH=$PATH:/home/test/cuda-10.1/bin/ //在 PATH 变量中加入需要使用的 cuda 版本的路径,使得系统可以使用 cuda 提供的可执行文件,包括 nvcc
想要永久设置上述 cuda 设置,用户可以直接在自己的 bash 设置文件 ~/.bashrc 文件尾部加入上述命令,保存后再通过 source ~/.bashrc 执行文件,即可完成当前终端的环境变量修改。如果需要使用新的 cuda 来编译文件,还可以通过 LD_LIBRARY_PATH 变量指定进行链接的 cuda 库文件的路径。
位于 ~/.bashrc 文件中的指令在每次终端启动时均会自动运行,后续本用户所打开的终端中的环境变量均会首先执行上述文件中的命令,从而获得对应的 cuda 变量。
>>>import torch >>>torch.version.cuda #输出一个 cuda 版本
如笔者环境下上述命令的输出如下图所示。
conda list | grep pytorch //查看安装的 Pytorch 的信息
笔者环境下上述命令的结果如图所示,可以看到显示的 cuda 信息与 torch.version.cuda 保持一致。
想要查看 Pytorch 实际使用的运行时的 cuda 目录,可以直接输出之前介绍的 cpp_extension.py 中的 CUDA_HOME 变量。
>>> import torch >>> import torch.utils >>> import torch.utils.cpp_extension >>> torch.utils.cpp_extension.CUDA_HOME #输出 Pytorch 运行时使用的 cuda