google-research

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

3
# Install tools and dependencies.
4
RUN apt-get -y update --fix-missing
5
RUN apt-get install -y \
6
	emacs \
7
	git \
8
	wget \
9
	libgoogle-glog-dev
10

11
# Install TensorFlow.
12
RUN apt-get install -y python3-dev python3-pip
13
RUN pip3 install --upgrade pip && pip install --upgrade tensorflow_gpu==2.3
14
RUN pip install --upgrade tensor2tensor
15

16
# Install CMake.
17
RUN apt-get install -y software-properties-common && \
18
    apt-get update && \
19
    wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | gpg --dearmor - | tee /etc/apt/trusted.gpg.d/kitware.gpg >/dev/null && \
20
    apt-add-repository 'deb https://apt.kitware.com/ubuntu/ bionic main' && \
21
    apt-get update && apt-get install -y cmake
22

23
# Install Sputnik.
24
RUN mkdir /mount
25
WORKDIR /mount
26
RUN git clone --recursive https://github.com/google-research/sputnik.git && \
27
	mkdir sputnik/build
28
WORKDIR /mount/sputnik/build
29
RUN cmake .. -DCMAKE_BUILD_TYPE=Release -DBUILD_TEST=OFF -DBUILD_BENCHMARK=OFF \
30
	-DCUDA_ARCHS="60;70" -DCMAKE_INSTALL_PREFIX=/usr/local/sputnik && \
31
	make -j8 install
32

33
# Copy the source into the image.
34
RUN mkdir -p /mount/sgk
35
COPY . /mount/sgk/
36

37
# Install SGK.
38
RUN mkdir /mount/sgk/build
39
WORKDIR /mount/sgk/build
40
RUN cmake .. -DCMAKE_BUILD_TYPE=Release -DCUDA_ARCHS="60;70" \
41
	-DCMAKE_INSTALL_PREFIX=/usr/local/sgk && \
42
	make -j8 install
43

44
# Setup the environment.
45
ENV PYTHONPATH="/mount:${PYTHONPATH}"
46
ENV LD_LIBRARY_PATH="`python3 -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())'`:${LD_LIBRARY_PATH}"
47
ENV LD_LIBRARY_PATH="/usr/local/sputnik/lib:${LD_LIBRARY_PATH}"
48

49
# Set the working directory.
50
WORKDIR /mount/sgk
51