dockerでpytorch 立ち上げ #memo

シンプルなtorch170のdockerの環境です。

環境：titanV
実行コマンド：nvidia-docker build --no-cache -t torch171_cu101 .

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

ARG use_cudnn=1

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get -y update
RUN apt-get -y install \
    sudo \
    apt-utils \
    git \
    curl \
    vim \
    unzip \
    wget \
    build-essential \
    cmake \
    libopenblas-dev \
    automake \
    openssh-client \
    openssh-server \
    lsof \
    python3.6 \
    python3-dev \
    python3-pip \
    python3-tk \
    libcupti-dev
RUN apt-get clean
RUN rm -rf /var/lib/apt/lists/*

#
# generate python script
#
RUN ln -s python3 /usr/bin/python
#
# update pip
#
RUN python3 -m pip install --upgrade pip==20.2.4 setuptools==50.3.2

# numpy 1.14.3  scipy 1.1.0
RUN python3 -m pip --no-cache-dir install numpy==1.14.3 scipy==1.1.0

RUN pip install torch==1.7.0+cu101 torchvision==0.8.1+cu101 torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html 

# RUN wget https://pypi.flyai.com/packages/c2/36/e9c43cfa694562954e7ae05d553f757b69beab6e2860faa46504a95ff489/nni-2.0-py3-none-manylinux1_x86_64.whl#sha256=fad8139253660e2ec72cfe2c8769e5c7bb4e909fea0baccfbfb088c6ff85f021
# RUN python3 -m pip install nni-2.0-py3-none-manylinux1_x86_64.whl

WORKDIR /root