Torch可识别GPU但TensorFlow无法识别的问题求助

你的问题核心是PyTorch和TensorFlow对CUDA依赖的处理方式不同，加上你用的python:3.9-slim轻量镜像本身不带CUDA工具链，导致TensorFlow无法识别GPU。

问题原因拆解 #

• PyTorch的pip包自带了CUDA Runtime组件，不需要系统提前安装CUDA Toolkit，只要Docker能通过nvidia-docker访问GPU，就能直接调用；
• TensorFlow则要求系统环境中必须安装对应版本的CUDA Toolkit和cuDNN，否则会判定没有GPU驱动；
• nvidia-smi显示的是宿主机驱动支持的最高CUDA版本（12.1），但容器内未安装CUDA Toolkit，所以nvcc找不到；python:slim是Debian轻量镜像，官方源没有NVIDIA的CUDA包，直接用apt安装不了。

解决方案 #

方案一：改用NVIDIA官方Python镜像（最省心）

NVIDIA提供了预装好CUDA、cuDNN和Python的镜像，直接基于这个构建，无需手动配置依赖。

修改后的Dockerfile：

ARG PYTHON_VERSION=3.9
# 选择适配TensorFlow的CUDA版本镜像，这里以CUDA 11.8为例（适配TensorFlow 2.13+）
FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04

# 安装Python 3.9并设置为默认
RUN apt-get update && apt-get install -y python3.9 python3-pip python3.9-dev
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.9 1

ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1

WORKDIR /app

RUN apt-get install -y ffmpeg

RUN --mount=type=cache,target=/root/.cache/pip \
    --mount=type=bind,source=requirements.txt,target=requirements.txt \
    python -m pip install -r requirements.txt
COPY . .
EXPOSE 80
CMD gunicorn 'main:app' --bind=0.0.0.0:80 --timeout=36000000 --workers=1 --threads=8

注意：要确保requirements.txt里的TensorFlow版本和CUDA版本匹配，比如TensorFlow 2.13.x对应CUDA 11.8，TensorFlow 2.15.x对应CUDA 12.2。

方案二：在python:slim中手动安装CUDA和cuDNN（适合必须用slim镜像的场景）

基于python:slim手动添加NVIDIA源，安装CUDA Runtime和cuDNN：

修改后的Dockerfile：

ARG PYTHON_VERSION=3.9
FROM python:${PYTHON_VERSION}-slim as base

ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1

WORKDIR /app

# 添加NVIDIA源并安装CUDA 11.8 Runtime和cuDNN
RUN apt-get update && apt-get install -y --no-install-recommends \
    gnupg2 wget && \
    wget https://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64/cuda-keyring_1.1-1_all.deb && \
    dpkg -i cuda-keyring_1.1-1_all.deb && \
    echo "deb https://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64/ /" > /etc/apt/sources.list.d/cuda.list && \
    apt-get update && \
    apt-get install -y --no-install-recommends \
    cuda-runtime-11-8 \
    libcudnn8=8.9.2.26-1+cuda11.8 && \
    # 清理缓存减小镜像体积
    apt-get clean && rm -rf /var/lib/apt/lists/* /cuda-keyring_1.1-1_all.deb

# 设置CUDA环境变量
ENV PATH=/usr/local/cuda-11.8/bin:$PATH
ENV LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH

RUN apt-get update && apt-get install -y ffmpeg

RUN --mount=type=cache,target=/root/.cache/pip \
    --mount=type=bind,source=requirements.txt,target=requirements.txt \
    python -m pip install -r requirements.txt
COPY . .
EXPOSE 80
CMD gunicorn 'main:app' --bind=0.0.0.0:80 --timeout=36000000 --workers=1 --threads=8

验证步骤 #

构建镜像并启动容器后，进入容器执行以下命令验证：

# 检查nvcc是否可用
nvcc --version
# 检查TensorFlow是否识别GPU
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

内容的提问来源于stack exchange，提问作者Fady's Cube