基于nvidia/cuda runtime镜像时PyTorch无法检测GPU的求助

问题描述 #

基于nvidia/cuda:11.3.1-cudnn8-runtime-ubuntu20.04构建Docker镜像后，启动容器执行以下代码返回False：

import torch
torch.cuda.is_available()

但切换到nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04基础镜像则正常返回True。需要在保留Runtime镜像体积优势的前提下修复PyTorch GPU检测问题。

原Dockerfile内容：

ARG CUDA_VERSION=11.3.1
FROM nvidia/cuda:${CUDA_VERSION}-cudnn8-runtime-ubuntu20.04
ARG PYTORCH_VERSION=1.12.1

# Set a docker label to enable container to use SAGEMAKER_BIND_TO_PORT environment variable if present
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
LABEL maintainer="Change Healthcare"
LABEL dlc_major_version="1"

ENV PATH /opt/conda/bin:$PATH

RUN rm /etc/apt/sources.list.d/*
RUN apt-get update
RUN apt-get install -y curl wget

RUN curl -L -o ~/miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-py38_23.1.0-1-Linux-x86_64.sh
RUN chmod +x ~/miniconda.sh
RUN ~/miniconda.sh -b -p /opt/conda
RUN rm ~/miniconda.sh
RUN /opt/conda/bin/conda install -y ruamel_yaml==0.15.100 cython botocore mkl-include mkl
RUN /opt/conda/bin/conda clean -ya

RUN pip install --upgrade pip --trusted-host pypi.org --trusted-host files.pythonhosted.org
RUN ln -s /opt/conda/bin/pip /usr/local/bin/pip
RUN ln -s /opt/conda/bin/pip /usr/local/bin/pip3
RUN ln -s /opt/conda/bin/python /usr/local/bin/python
RUN pip install packaging==20.4 enum-compat==0.0.3

# Conda installs links for libtinfo.so.6 and libtinfo.so.6.2 both
# Which causes "/opt/conda/lib/libtinfo.so.6: no version information available" warning
# Removing link for libtinfo.so.6. This change is needed only for ubuntu 20.04-conda, and can be reverted
# once conda fixes the issue: https://github.com/conda/conda/issues/9680
RUN rm -rf /opt/conda/lib/libtinfo.so.6

WORKDIR /

RUN cd tmp/ \
 && rm -rf tmp*

# Uninstall and re-install torch and torchvision from the PyTorch website
RUN pip uninstall -y torch
RUN /opt/conda/bin/conda install pytorch==${PYTORCH_VERSION} cudatoolkit=11.3 -c pytorch

解决方案 #

1. 修改Dockerfile，调整PyTorch安装方式 #

将原Dockerfile中通过conda安装PyTorch的步骤替换为pip安装预编译的CUDA兼容版本，同时添加CUDA库路径环境变量：

ARG CUDA_VERSION=11.3.1
FROM nvidia/cuda:${CUDA_VERSION}-cudnn8-runtime-ubuntu20.04
ARG PYTORCH_VERSION=1.12.1

# Set a docker label to enable container to use SAGEMAKER_BIND_TO_PORT environment variable if present
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
LABEL maintainer="Change Healthcare"
LABEL dlc_major_version="1"

# 配置CUDA库路径，确保PyTorch能找到GPU驱动依赖
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
ENV PATH /opt/conda/bin:$PATH

RUN rm /etc/apt/sources.list.d/*
RUN apt-get update
RUN apt-get install -y curl wget

RUN curl -L -o ~/miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-py38_23.1.0-1-Linux-x86_64.sh
RUN chmod +x ~/miniconda.sh
RUN ~/miniconda.sh -b -p /opt/conda
RUN rm ~/miniconda.sh
RUN /opt/conda/bin/conda install -y ruamel_yaml==0.15.100 cython botocore mkl-include mkl
RUN /opt/conda/bin/conda clean -ya

RUN pip install --upgrade pip --trusted-host pypi.org --trusted-host files.pythonhosted.org
RUN ln -s /opt/conda/bin/pip /usr/local/bin/pip
RUN ln -s /opt/conda/bin/pip /usr/local/bin/pip3
RUN ln -s /opt/conda/bin/python /usr/local/bin/python
RUN pip install packaging==20.4 enum-compat==0.0.3

# Conda installs links for libtinfo.so.6 and libtinfo.so.6.2 both
# Which causes "/opt/conda/lib/libtinfo.so.6: no version information available" warning
# Removing link for libtinfo.so.6. This change is needed only for ubuntu 20.04-conda, and can be reverted
# once conda fixes the issue: https://github.com/conda/conda/issues/9680
RUN rm -rf /opt/conda/lib/libtinfo.so.6

WORKDIR /

RUN cd tmp/ \
 && rm -rf tmp*

# 改用pip安装PyTorch官方预编译的CUDA 11.3兼容包
RUN pip uninstall -y torch
RUN pip install torch==${PYTORCH_VERSION}+cu113 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu113

2. 正确启动容器 #

启动容器时必须使用NVIDIA容器运行时，确保GPU设备被正确挂载：

docker run --gpus all -it your-built-image-name bash

3. 验证GPU检测 #

进入容器后执行以下代码，确认PyTorch能检测到GPU：

import torch
print(torch.cuda.is_available())  # 预期输出: True
print(torch.cuda.device_count())  # 输出可用GPU数量

原因说明 #

• CUDA Runtime镜像仅包含运行CUDA程序所需的最小依赖库，缺少conda安装PyTorch时依赖的部分CUDA开发组件（如头文件、静态链接库），导致PyTorch无法完成CUDA环境初始化。
• PyTorch官方提供的pip预编译包已内置所需的CUDA Runtime组件，无需依赖系统的CUDA开发库，因此可以在Runtime镜像的轻量环境中正常识别GPU。

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