Docker中使用PyTorch镜像无法设置CUDA_HOME环境变量的求助

问题原因 #

你使用的pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime属于runtime版本镜像，仅包含运行PyTorch所需的CUDA runtime库，没有完整的CUDA toolkit安装目录，所以/usr/local/下看不到cuda文件夹，这是正常设计。

解决方案 #

根据你的需求（只需运行PyTorch，无需编译自定义CUDA扩展），有以下几种处理方式：

1. 临时设置CUDA_HOME（单次会话生效）

利用PyTorch内置的CUDA路径变量直接设置：

export CUDA_HOME=$(python -c "import torch; print(torch.utils.cpp_extension.CUDA_HOME)")

验证是否生效：

echo $CUDA_HOME

2. 永久设置CUDA_HOME（容器启动自动生效）

• 若基于该镜像构建自定义镜像，在Dockerfile中添加：

  RUN echo 'export CUDA_HOME=$(python -c "import torch; print(torch.utils.cpp_extension.CUDA_HOME)")' >> ~/.bashrc
  

• 若直接启动容器，通过环境变量参数传递：

  docker run -it --gpus all -e CUDA_HOME=$(python -c "import torch; print(torch.utils.cpp_extension.CUDA_HOME)") pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime
  

3. 切换到devel版本镜像（需完整CUDA toolkit时使用）

如果需要编译自定义CUDA扩展等场景，改用同版本的devel镜像，该镜像包含完整CUDA toolkit：

docker pull pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel

进入容器后直接设置：

export CUDA_HOME=/usr/local/cuda

注意事项 #

你的需求是PyTorch 1.11+版本高于9.2的CUDA，当前runtime镜像已满足运行PyTorch的核心需求，仅当需要编译CUDA相关扩展时，才需要切换到devel版本镜像。

内容的提问来源于stack exchange，提问作者tarashi