Python env激活时自动选用合适CUDA版本

Intro

在使用不同的深度学习框架(TensorFlow, Caffe, PyTorch, Caffe2等),甚至同一个框架的不同版本时,可能要求的CUDA和CuDNN的版本不尽相同.如目前最新的TF和PyTorch分别是1.6和0.32版本,要求的CUDA版本分别是9.0和9.1,而实验室公用机器默认安装的是CUDA8.0,为了最小化影响公用机器,需要找到一种方法可以使得多版本的CUDA共存,并不会互相影响.
根据之前CUDA的安装经验,CUDA起作用是通过系统变量的PATHLD_LIBRARY_PATH,那么可以尝试将CUDA安装到Home或其他用户目录中而非默认的/usr/local/cuda中,然后通过合理设置系统变量使CUDA生效.
本文中系统的基本参数为: - Ubuntu 16.04 x64 - miniconda

CUDA安装

到NVIDIA(NV)官网下载相应的CUDA安装包和CuDNN的安装包,CUDA安装包通常包含一个基础包和若干补丁包(我通常用run,不用deb包),补丁包的命名为基础包的版本号后加.x(x为数字)表示第几个补丁,安装时,按补丁顺序依次安装.CuDNN选择linux通用的tar.gz的压缩包即可.
给run文件可执行权限,可以右键在权限中设置,也可以chmod +x 文件名设置.在bash中执行run文件,注意阅读提示. 印象中 第一个是是否同意用户协议,输入accept;第二个问是否要装驱动,输入n,第三个问安装位置,默认在/usr/local/cuda,如安装到该位置需要root权限,可以更改到用户自己的目录中,输入需要安装的目录,如~/cuda/cuda9.0;再下一个是是否创建link,也是要root权限的,选不要创建;最后开始将CUDA安装到指定目录.然后再依次安装补丁包,补丁包选项比较少,只有同意协议和选择CUDA安装位置,注意不要输错路径就行.
解压下载的CuDNN安装包,将cuda目录下的binlib64复制相应的CUDA目录下,即完成安装.

Python配置

以PyTorch为例,选择安装基于CUDA9.1和Python3.6的PyTorch.

  • 创建env
    conda create -n torch python=3.6

  • 安装PyTorch
    参考官网的提示,选择用pip方式安装,因为conda提供的包比较旧

    1
    2
    pip install http://download.pytorch.org/whl/cu91/torch-0.3.1-cp36-cp36m-linux_x86_64.whl
    pip install torchvision

  • 设置activate和deactivate对环境变量的修改 参考文档,进入env的所在目录,新建两个文件夹和两个sh文件.

    1
    2
    3
    4
    5
    cd ~/miniconda/envs/torch
    mkdir -p ./etc/conda/activate.d
    mkdir -p ./etc/conda/deactivate.d
    touch ./etc/conda/activate.d/env_vars.sh
    touch ./etc/conda/deactivate.d/env_vars.sh

修改activate.denv_vars.sh的内容为:

1
2
3
4
5
6
7
#!/bin/sh

export OLD_PATH=${PATH}
export OLD_LD_LIBRARY_PATH=${LD_LIBRARY_PATH}

export PATH=/home/username/cuda/cuda-9.1/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/home/username/cuda/cuda-9.1/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

修改deactivate.denv_vars.sh的内容为:

1
2
3
4
5
6
7
#!/bin/sh

unset PATH
unset LD_LIBRARY_PATH

export PATH=${OLD_PATH}
export LD_LIBRARY_PATH=${OLD_LD_LIBRARY_PATH}

注意:修改上面的cuda路径为自己的!

测试效果

1
2
3
4
5
6
7
8
9
10
echo $PATH
echo $LD_LIBRARY_PATH

source activate torch
echo $PATH
echo $LD_LIBRARY_PATH

source deactivate torch
echo $PATH
echo $LD_LIBRARY_PATH

观察输出路径的不同.