Ubuntu镜像中配置带GPU/CUDA的TensorFlow 2.14 Docker镜像遇阻求助

Dockerfile中的错误点 #

• Conda环境激活逻辑失效：每个RUN指令都是独立的shell会话，在同一个RUN里执行conda init && . ~/.bashrc后，后续的conda activate tf并不会真正激活环境，导致TensorFlow被安装到base环境而非目标tf环境。而且conda init在构建阶段的非交互式shell中作用有限，完全没必要反复执行。
• SHELL与CMD指令冲突：已经通过SHELL指令指定了用conda run启动tf环境的shell，但CMD里又重复执行conda init和conda activate，反而触发环境切换错误。conda run本身会自动激活指定环境，无需手动操作。
• 环境变量配置缺失：仅设置了Miniconda的PATH，但CUDA相关库的路径没有加入LD_LIBRARY_PATH，导致TensorFlow运行时找不到CUDA和cuDNN库。
• 镜像冗余未清理：Miniconda安装包下载后未删除，额外增加了镜像体积。

TensorFlow无法识别GPU的原因 #

• 容器未正确挂载GPU驱动：虽然WSL2主机有可用GPU，但Docker容器默认不会自动挂载NVIDIA驱动文件，必须通过--gpus all参数启动容器，同时主机要安装nvidia-docker2工具来支持GPU容器。
• CUDA环境变量未加载：Conda安装的cudatoolkit和cudnn的库路径没有加入系统环境变量，TensorFlow无法定位到这些依赖库。
• 版本兼容性隐患：TensorFlow 2.14官方推荐搭配CUDA 11.8和cuDNN 8.6，你使用的CUDA 11.5虽然理论上兼容，但版本匹配度不如官方推荐组合，可能存在底层库适配问题。

修正后的Dockerfile #

FROM ubuntu:22.04
LABEL maintainer="test-user"

# 安装依赖并完成Miniconda安装，清理冗余文件
RUN apt-get update && \
    apt-get install -y wget && \
    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
    bash Miniconda3-latest-Linux-x86_64.sh -b && \
    rm Miniconda3-latest-Linux-x86_64.sh && \
    rm -rf /var/lib/apt/lists/*

# 配置Conda路径并创建环境，直接用conda run安装依赖
ENV PATH="/root/miniconda3/bin:$PATH"
RUN conda create --name tf python=3.11.0 -y && \
    conda run -n tf conda install -c conda-forge cudatoolkit=11.5 cudnn=8.9.7 -y && \
    conda run -n tf pip install tensorflow==2.14.0

# 设置CUDA库路径和NVIDIA容器相关环境变量
ENV LD_LIBRARY_PATH="/root/miniconda3/envs/tf/lib:$LD_LIBRARY_PATH"
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility

# 复制检测脚本
COPY gpu_check.py /root/gpu_check.py

# 直接用conda run启动环境并执行检测脚本
CMD ["conda", "run", "-n", "tf", "python", "/root/gpu_check.py"]

运行注意事项 #

1. 确保主机已安装nvidia-docker2并重启Docker服务：

   sudo apt-get install nvidia-docker2
   sudo systemctl restart docker
   

2. 启动容器时必须添加--gpus all参数：

   docker run --gpus all tf-test
   

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