构建Vision Transformer Dockerfile时遇jax[gpu]>=0.3.4版本匹配错误

构建Vision Transformer的Dockerfile时遇到如下依赖安装错误：

ERROR: Could not find a version that satisfies the requirement
jax[gpu]>=0.3.4 (from -r vit_jax/requirements.txt (line 8)) (from
versions: 0.0, 0.1, 0.1.1, ..., 0.2.17)
ERROR: No matching distribution found for jax[gpu]>=0.3.4 (from -r vit_jax/requirements.txt (line 8))

使用的Dockerfile内容如下：

FROM pytorch/pytorch:1.2-cuda10.0-cudnn7-runtime

ENV DEBIAN_FRONTEND=noninteractive
ARG USERNAME=user
WORKDIR /dockertest
ARG WORKDIR=/dockertest
RUN apt-get update && apt-get install -y \
        automake autoconf libpng-dev nano python3-pip \
        sudo curl zip unzip libtool swig zlib1g-dev pkg-config \
        python3-mock libpython3-dev libpython3-all-dev \
        g++ gcc cmake make pciutils cpio gosu wget \
        libgtk-3-dev libxtst-dev sudo apt-transport-https \
        build-essential gnupg git xz-utils vim libgtk2.0-0 libcanberra-gtk-module\
        libva-dev libdrm-dev xorg xorg-dev protobuf-compiler \
        openbox libx11-dev libgl1-mesa-glx libgl1-mesa-dev \
        libtbb2 libtbb-dev libopenblas-dev libopenmpi-dev \
    && sed -i 's/# set linenumbers/set linenumbers/g' /etc/nanorc \
    && apt clean \
    && rm -rf /var/lib/apt/lists/*

RUN git clone https://github.com/google-research/vision_transformer.git \
    &&cd vision_transformer \
    && pip3 install pip --upgrade \
    && pip install -r vit_jax/requirements.txt \
    &&python -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \
    --config=$(pwd)/vit_jax/configs/vit.py:b16,cifar10 \
    --config.pretrained_dir='gs://vit_models/imagenet21k' \
    && pip cache purge

RUN echo "root:root" | chpasswd \
    && adduser --disabled-password --gecos "" "${USERNAME}" \
    && echo "${USERNAME}:${USERNAME}" | chpasswd \
    && echo "%${USERNAME}    ALL=(ALL)   NOPASSWD:    ALL" >> /etc/sudoers.d/${USERNAME} \
    && chmod 0440 /etc/sudoers.d/${USERNAME}
USER ${USERNAME}
RUN sudo chown -R ${USERNAME}:${USERNAME} ${WORKDIR}
WORKDIR ${WORKDIR}

核心原因 #

你使用的基础镜像pytorch/pytorch:1.2-cuda10.0-cudnn7-runtime过于老旧：

• 该镜像基于Python 3.6，而jax从0.3.0版本开始不再支持Python 3.6，因此pip只能找到0.2.x及更早的jax版本，无法满足jax[gpu]>=0.3.4的要求。
• 同时，CUDA 10.0也不在jax 0.3+版本的支持范围内，jax后续版本需要CUDA 11.x及以上。

修改方案 #

1. 更换基础镜像 #

选择支持Python 3.7+、CUDA 11.x的PyTorch镜像，比如：

FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime

这个镜像满足Python版本和CUDA版本要求，可兼容jax 0.3+的安装。

2. 优化Dockerfile步骤 #

• 拆分依赖安装和代码克隆步骤，利用Docker缓存提升构建效率。
• 明确使用pip3避免Python版本混淆。
• 移除不必要的系统依赖，减少镜像体积。

修改后的Dockerfile示例：

FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime

ENV DEBIAN_FRONTEND=noninteractive
ARG USERNAME=user
ARG WORKDIR=/dockertest
WORKDIR ${WORKDIR}

# 安装必要系统依赖
RUN apt-get update && apt-get install -y --no-install-recommends \
        nano git vim sudo curl build-essential \
    && sed -i 's/# set linenumbers/set linenumbers/g' /etc/nanorc \
    && apt clean \
    && rm -rf /var/lib/apt/lists/*

# 克隆Vision Transformer代码并安装依赖
RUN git clone https://github.com/google-research/vision_transformer.git \
    && cd vision_transformer \
    && pip3 install --upgrade pip \
    && pip3 install -r vit_jax/requirements.txt \
    && pip3 cache purge

# 执行测试命令（可选，验证安装成功）
RUN cd vision_transformer \
    && python3 -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \
    --config=$(pwd)/vit_jax/configs/vit.py:b16,cifar10 \
    --config.pretrained_dir='gs://vit_models/imagenet21k'

# 创建非root用户
RUN echo "root:root" | chpasswd \
    && adduser --disabled-password --gecos "" "${USERNAME}" \
    && echo "${USERNAME}:${USERNAME}" | chpasswd \
    && echo "%${USERNAME}    ALL=(ALL)   NOPASSWD:    ALL" >> /etc/sudoers.d/${USERNAME} \
    && chmod 0440 /etc/sudoers.d/${USERNAME} \
    && chown -R ${USERNAME}:${USERNAME} ${WORKDIR}

USER ${USERNAME}
WORKDIR ${WORKDIR}

3. 额外注意事项 #

• 如果不需要PyTorch，也可以直接使用NVIDIA的CUDA基础镜像搭配Python，比如nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu22.04，然后手动安装Python和依赖。
• 若需要保留旧版本CUDA，可修改vit_jax/requirements.txt中的jax版本限制为jax[gpu]<=0.2.17，但这会导致无法使用jax后续的新特性，不推荐。

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