有没有办法检查 CUDA 是否可用并设置环境变量（如果可用）？

Question

我有一个用于机器学习的 Python 程序。我想使用 Shell 脚本来运行它，有时我会设置环境变量CUDA_VISIBLE_DEVICES以便使用 NVIDIA GPU。

我想在没有安装 CUDA 的 M1 MacBook 上运行相同的程序。有没有办法根据是否安装CUDA来设置环境变量？

我检查 CUDA 是否可用（在 Linux 中）的常用方法是nvcc --version.如果返回有效的输出，则它已安装。我在想这样的事情：

if "nvcc --version":
    export CUDA_VISIBLE_DEVICES=0

Answer 1

在 bash 脚本中，语法更像是这样：

if nvcc --version
then
  export CUDA_VISIBLE_DEVICES=0
fi

或者，更短一点：

nvcc --version && export CUDA_VISIBLE_DEVICES=0

这两者都测试nvcc --version命令的退出状态，而不检查其输出的有效性。