File size: 1,533 Bytes
3c849be |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
FROM nvidia/cuda:11.0-base-ubuntu20.04
ENV DEBIAN_FRONTEND noninteractive
ENV CUDNN_VERSION=8.0.5.39-1+cuda11.1
ENV NCCL_VERSION=2.7.8-1+cuda11.1
ARG python=3.8
ENV PYTHON_VERSION=${python}
# Set default shell to /bin/bash
SHELL ["/bin/bash", "-cu"]
RUN apt-get update && apt-get install -y --allow-downgrades \
--allow-change-held-packages --no-install-recommends \
build-essential \
cmake \
git \
curl \
vim \
wget \
ca-certificates \
libcudnn8=${CUDNN_VERSION} \
libnccl2=${NCCL_VERSION} \
libnccl-dev=${NCCL_VERSION} \
libjpeg-dev \
libpng-dev \
python${PYTHON_VERSION} \
python${PYTHON_VERSION}-dev \
python${PYTHON_VERSION}-distutils \
librdmacm1 \
libibverbs1 \
ibverbs-providers
RUN ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python
RUN curl -O https://bootstrap.pypa.io/get-pip.py && \
python get-pip.py && \
rm get-pip.py
RUN /usr/bin/python -m pip install --upgrade pip
# Install pytorch
RUN pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 \
-f https://download.pytorch.org/whl/torch_stable.html
RUN pip install tensorboard==2.5.0
RUN pip install tensorboard-data-server==0.6.1
RUN pip install tensorboard-plugin-wit==1.8.0
RUN pip install tensorboardX==1.8
RUN pip install timm==0.4.5
RUN pip install opencv-contrib-python-headless==4.5.2.54
RUN pip install tqdm==4.61.2
RUN pip install PyYAML==5.4.1
RUN pip install Pillow==8.3.1
RUN pip install einops==0.3.0
RUN pip install scipy==1.7.1
|