其他受访者已经描述了可以使用哪些命令来检查 CUDA 版本。在这里,我将描述如何将这些命令的输出转换为“10.2”、“11.0”等形式的环境变量。
回顾一下,你可以使用
nvcc --version
查找 CUDA 版本。
我认为这应该是您的第一个停靠港。
如果您安装了多个版本的 CUDA,此命令应打印出 PATH 上最高副本的版本。
输出如下所示:
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Thu_Jun_11_22:26:38_PDT_2020
Cuda compilation tools, release 11.0, V11.0.194
Build cuda_11.0_bu.TC445_37.28540450_0
我们可以通过 sed 传递这个输出来选择 MAJOR.MINOR 发布版本号。
CUDA_VERSION=$(nvcc --version | sed -n 's/^.*release \([0-9\.]\+\).*$/\1/p')
如果 nvcc 不在您的路径上,您应该能够通过指定 nvcc 的默认位置的完整路径来运行它。
/usr/local/cuda/bin/nvcc --version
其中的输出同上,同样可以解析。
或者,您可以从 version.txt 文件中找到 CUDA 版本。
cat /usr/local/cuda/version.txt
其中的输出
CUDA Version 10.1.243
可以使用 sed 进行解析,以找出 MAJOR.MINOR 发布版本号。
CUDA_VERSION=$(cat /usr/local/cuda/version.txt | sed 's/.* \([0-9]\+\.[0-9]\+\).*/\1/')
请注意,有时 version.txt 文件指的是与 nvcc --version 不同的 CUDA 安装。在这种情况下,nvcc 版本应该是您实际使用的版本。
我们可以将这三种方法结合在一起,以稳健地获得 CUDA 版本,如下所示:
if nvcc --version 2&> /dev/null; then
# Determine CUDA version using default nvcc binary
CUDA_VERSION=$(nvcc --version | sed -n 's/^.*release \([0-9\.]\+\).*$/\1/p');
elif /usr/local/cuda/bin/nvcc --version 2&> /dev/null; then
# Determine CUDA version using /usr/local/cuda/bin/nvcc binary
CUDA_VERSION=$(/usr/local/cuda/bin/nvcc --version | sed -n 's/^.*release \([0-9\.]\+\).*$/\1/p');
elif [ -f "/usr/local/cuda/version.txt" ]; then
# Determine CUDA version using /usr/local/cuda/version.txt file
CUDA_VERSION=$(cat /usr/local/cuda/version.txt | sed 's/.* \([0-9]\+\.[0-9]\+\).*/\1/')
else
CUDA_VERSION=""
fi
此环境变量对于下游安装很有用,例如当 pip 安装为正确 CUDA 版本编译的 pytorch 副本时。
python -m pip install \
"torch==1.9.0+cu${CUDA_VERSION/./}" \
"torchvision==0.10.0+cu${CUDA_VERSION/./}" \
-f https://download.pytorch.org/whl/torch_stable.html
同样,你可以在没有安装 CUDA 的情况下安装 CPU 版本的 pytorch。
if [ "$CUDA_VERSION" = "" ]; then
MOD="+cpu";
echo "Warning: Installing CPU-only version of pytorch"
else
MOD="+cu${CUDA_VERSION/./}";
echo "Installing pytorch with $MOD"
fi
python -m pip install \
"torch==1.9.0${MOD}" \
"torchvision==0.10.0${MOD}" \
-f https://download.pytorch.org/whl/torch_stable.html
但请注意这一点,因为当您打算获得 GPU 支持时,您可能会意外安装仅 CPU 版本。
例如,如果您在没有 GPU 的服务器登录节点上运行安装脚本,您的作业将部署到有 GPU 的节点上。在这种情况下,登录节点通常不会安装 CUDA。