Dockerfile.pytorch

FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04

ARG PYTORCH_VER
ARG ONNX_VER=1.6.0

###########################################
## update & install packages
###########################################
SHELL ["/bin/bash", "-c"]

RUN apt-get -qq update && \
    apt-get -qq install -y curl wget vim htop git software-properties-common unzip zip unrar rar cmake gfortran && \
    apt-get -qq install -y build-essential libfreetype6-dev libhdf5-serial-dev libzmq3-dev pkg-config libsm6 libxext6 libxrender-dev

###########################################
## Install python 3.6
###########################################
RUN add-apt-repository ppa:deadsnakes/ppa && \
    apt-get -qq update && \
    apt-get -qq install -y build-essential python3.6 python3.6-dev python3-pip && \
    python3.6 -m pip install -U pip && \
    python3.6 -m pip install wheel
RUN ln -s `which python3.6` /usr/bin/python

###########################################
## Install pytorch
###########################################
RUN python3.6 -m pip install torch==${PYTORCH_VER} && \
    python3.6 -m pip install onnx==${ONNX_VER} && \
    python3.6 -m pip install numpy