-
-
Save ruipeterpan/f64ea87e5974b15a3b07bc93e50b0719 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
FROM nvidia/cuda:11.0-devel-ubuntu18.04 | |
ARG https_proxy | |
ARG http_proxy | |
ARG BYTEPS_BASE_PATH=/usr/local | |
ARG BYTEPS_PATH=$BYTEPS_BASE_PATH/byteps | |
ARG BYTEPS_GIT_LINK=https://github.com/bytedance/byteps | |
ARG BYTEPS_BRANCH=master | |
ARG DEBIAN_FRONTEND=noninteractive | |
RUN apt-get update | |
RUN apt-get install -y --allow-downgrades --allow-change-held-packages --no-install-recommends \ | |
build-essential \ | |
tzdata \ | |
ca-certificates \ | |
git \ | |
curl \ | |
wget \ | |
vim \ | |
gdb \ | |
cmake \ | |
lsb-release \ | |
libcudnn8=8.0.5.39-1+cuda11.0 \ | |
libnuma-dev \ | |
ibverbs-providers \ | |
librdmacm-dev \ | |
ibverbs-utils \ | |
rdmacm-utils \ | |
libibverbs-dev \ | |
python3 \ | |
python3-dev \ | |
python3-pip \ | |
python3-setuptools \ | |
libnccl2=2.7.8-1+cuda11.0 \ | |
libnccl-dev=2.7.8-1+cuda11.0 | |
# install framework | |
# note: for tf <= 1.14, you need gcc-4.9 | |
ARG FRAMEWORK=pytorch | |
RUN if [ "$FRAMEWORK" = "tensorflow" ]; then \ | |
pip3 install --upgrade pip; \ | |
pip3 install -U tensorflow-gpu==1.15.0; \ | |
elif [ "$FRAMEWORK" = "pytorch" ]; then \ | |
python3 -m pip install torch==1.7.0+cu110 torchvision==0.8.1+cu110 -f https://download.pytorch.org/whl/torch_stable.html; \ | |
elif [ "$FRAMEWORK" = "mxnet" ]; then \ | |
pip3 install -U mxnet-cu100==1.5.0; \ | |
else \ | |
echo "unknown framework: $FRAMEWORK"; \ | |
exit 1; \ | |
fi | |
ENV LD_LIBRARY_PATH /usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH | |
RUN cd $BYTEPS_BASE_PATH &&\ | |
git clone --recursive -b $BYTEPS_BRANCH $BYTEPS_GIT_LINK &&\ | |
cd $BYTEPS_PATH &&\ | |
python3 setup.py install |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment