Create dockers for nestedtensor pytorch extension for varying size tensors

nested-tensor-cpu.dockerfile

FROM ubuntu:18.04
# copies https://raw.githubusercontent.com/pytorch/pytorch/master/docker/pytorch/Dockerfile
ARG PYTHON_VERSION=3.6.10
ARG WITH_TORCHVISION=1
RUN apt-get update && apt-get install -y --no-install-recommends \
         build-essential \
         cmake \
         git \
         curl \
         ca-certificates \
         libjpeg-dev \
         libpng-dev \
	 openssh-server && \
     rm -rf /var/lib/apt/lists/*


RUN curl -o ~/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
     chmod +x ~/miniconda.sh && \
     ~/miniconda.sh -b -p /opt/conda && \
     rm ~/miniconda.sh && \
     /opt/conda/bin/conda install -y python=$PYTHON_VERSION numpy pyyaml scipy ipython mkl mkl-include ninja cython typing && \
     /opt/conda/bin/conda install -y -c pytorch-nightly pytorch cpuonly && \
     /opt/conda/bin/conda install -y -c pytorch-nightly ipython notebook && \
     /opt/conda/bin/conda clean -ya
ENV PATH /opt/conda/bin:$PATH
# This must be done before pip so that requirements.txt is available
WORKDIR /opt/pytorch
COPY . .

ENV TORCH_CUDA_ARCH_LIST "3.5 5.2 6.0 6.1 7.0+PTX" 
ENV TORCH_NVCC_FLAGS "-Xfatbin -compress-all" 
ENV CMAKE_PREFIX_PATH "$(dirname $(which conda))/../" 
#    pip install -v .

RUN if [ "$WITH_TORCHVISION" = "1" ] ; then /opt/conda/bin/conda install -y torchvision -c pytorch-nightly  && /opt/conda/bin/conda; fi

RUN git clone https://github.com/pytorch/nestedtensor.git && cd nestedtensor && pip install -v .

WORKDIR /workspace
COPY ./nested_simple_check.py .
RUN chmod -R a+w .
RUN chmod +x nested_simple_check.py . && ./nested_simple_check.py --no-cuda

nested-tensor.dockerfile

FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04
# copies https://raw.githubusercontent.com/pytorch/pytorch/master/docker/pytorch/Dockerfile
ARG PYTHON_VERSION=3.6.10
ARG WITH_TORCHVISION=1
RUN apt-get update && apt-get install -y --no-install-recommends \
         build-essential \
         cmake \
         git \
         curl \
         ca-certificates \
         libjpeg-dev \
         libpng-dev \
	 openssh-server && \
     rm -rf /var/lib/apt/lists/*


RUN curl -o ~/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
     chmod +x ~/miniconda.sh && \
     ~/miniconda.sh -b -p /opt/conda && \
     rm ~/miniconda.sh && \
     /opt/conda/bin/conda install -y python=$PYTHON_VERSION numpy pyyaml scipy ipython mkl mkl-include ninja cython typing && \
     /opt/conda/bin/conda install -y -c pytorch magma-cuda102 && \
     /opt/conda/bin/conda install -y -c pytorch-nightly pytorch && \
     /opt/conda/bin/conda install -y -c pytorch-nightly ipython notebook && \
     /opt/conda/bin/conda clean -ya
ENV PATH /opt/conda/bin:$PATH
# This must be done before pip so that requirements.txt is available
WORKDIR /opt/pytorch
COPY . .

ENV TORCH_CUDA_ARCH_LIST "3.5 5.2 6.0 6.1 7.0+PTX" 
ENV TORCH_NVCC_FLAGS "-Xfatbin -compress-all" 
ENV CMAKE_PREFIX_PATH "$(dirname $(which conda))/../" 

RUN if [ "$WITH_TORCHVISION" = "1" ] ; then conda install -y torchvision -c pytorch-nightly  && /opt/conda/bin/conda clean -ya; fi

RUN git clone https://github.com/pytorch/nestedtensor.git && cd nestedtensor && pip install -v .

WORKDIR /workspace
COPY ./nested_simple_check.py .
RUN chmod -R a+w .
RUN chmod +x nested_simple_check.py .  &&\
    ./nested_simple_check.py --no-cuda &&\
    echo "Run to check with cuda:\n" && echo "docker run --gpus all -it nestedtensor ./nested_simple_check.py"
# Can't run it with cuda since the builder runs without cuda && ./nested_simple_check.py

nested_simple_check.py

#!/opt/conda/bin/python
from argparse import ArgumentParser
import sys
import nestedtensor
from nestedtensor import tensorwise
from nestedtensor import torch


def random_nested():
    return nestedtensor.nested_tensor([
                                       [
                                        torch.rand(2, 3),
                                        torch.rand(4, 5)
                                       ],
                                       [
                                        torch.rand(1, 2)
                                       ]
                                     ])


@tensorwise()
def simple_fn(t1, t2):
    return t1 + 1 + t2


if __name__ == '__main__':
    parser = ArgumentParser('A small check of the pytorch '
                            'nested tensor library')
    parser.add_argument('--cuda',
                        action='store_true',
                        dest='cuda')
    parser.add_argument('--no-cuda',
                        action='store_false',
                        dest='cuda')
    parser.set_defaults(cuda=True)
    args = parser.parse_args()
    try:
        t1 = random_nested()
        t2 = random_nested()
        if args.cuda == True:
            device = 'cuda'
            t1.to(device)
            t2.to(device)
        simple_fn(t1, t1)
        print('Everything seems to be fine')
    except Exception as e:
        print('Error encountered')
        print(e)
        sys.exit(1)