Ubuntu20.04构建Vision Transformer Docker镜像时遇libcudart.so.11.0缺失错误

问题 #

在Ubuntu 20.04系统上构建Vision Transformer的Dockerfile时，出现两个问题：

1. 无法加载动态库libcudart.so.11.0，提示文件不存在
2. 构建过程在执行git clone相关步骤时终止

错误详情：

2022-11-04 09:08:49.205922: W external/org_tensorflow/tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64]
Could not load dynamic library 'libcudart.so.11.0'; dlerror:
libcudart.so.11.0: cannot open shared object file: No such file or
directory; LD_LIBRARY_PATH:
/usr/local/nvidia/lib:/usr/local/nvidia/lib64

使用的Dockerfile内容：

FROM pytorch/pytorch:1.8.1-cuda10.2-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-drm2 libva-x11-2 vainfo libva-wayland2 libva-glx2 \
        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}

解决方案 #

1. 解决libcudart.so.11.0缺失问题

• 问题根源：当前基础镜像pytorch/pytorch:1.8.1-cuda10.2-cudnn7-runtime自带CUDA 10.2，而依赖库需要CUDA 11.0版本的libcudart.so，版本不匹配导致文件缺失。
• 修复方法：更换基础镜像为对应CUDA 11.x版本的PyTorch镜像，比如pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime，确保镜像包含libcudart.so.11.0。

2. 修复git clone步骤终止问题

• 问题根源：构建阶段执行python -m vit_jax.main会直接启动训练，构建环境通常无GPU资源，且该操作会导致构建过程卡住或失败。
• 修复方法：移除构建阶段的训练启动命令，改为容器启动后手动执行。修改后的RUN指令如下：

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 \
    && pip cache purge

优化后的完整Dockerfile

FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-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 apt-transport-https \
        build-essential gnupg git xz-utils vim libgtk2.0-0 libcanberra-gtk-module\
        # libva-drm2 libva-x11-2 vainfo libva-wayland2 libva-glx2 \
        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 \
    && 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}

额外提示

• 构建镜像前确保Docker环境已安装nvidia-docker，支持GPU资源调用。
• 可精简apt-get install的包列表，移除不需要的依赖，减少镜像体积。

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