ARG PYTORCH="1.9.0" ARG CUDA="11.1" ARG CUDNN="8" FROM pytorch/pytorch:${PYTORCH}-cuda${CUDA}-cudnn${CUDNN}-devel ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6+PTX" \ TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \ CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \ FORCE_CUDA="1" RUN rm /etc/apt/sources.list.d/cuda.list \ && rm /etc/apt/sources.list.d/nvidia-ml.list \ && apt-key del 7fa2af80 \ && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub \ && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub # (Optional) # RUN sed -i 's/http:\/\/archive.ubuntu.com\/ubuntu\//http:\/\/mirrors.aliyun.com\/ubuntu\//g' /etc/apt/sources.list && \ # pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple RUN apt-get update \ && apt-get install -y ffmpeg libsm6 libxext6 git ninja-build libglib2.0-0 libxrender-dev \ && apt-get clean \ && rm -rf /var/lib/apt/lists/* # Install MMEngine , MMCV and MMDet RUN pip install --no-cache-dir openmim && \ mim install --no-cache-dir "mmengine>=0.6.0" "mmcv>=2.0.0rc4,<2.1.0" "mmdet>=3.0.0,<4.0.0" # Install MMYOLO RUN git clone https://github.com/open-mmlab/mmyolo.git /mmyolo && \ cd /mmyolo && \ mim install --no-cache-dir -e . WORKDIR /mmyolo