How to use GPU with stack, docker, nvidia-docker (#157)

README.md

@@ -82,6 +82,19 @@ There is also a demo application:
     cd tensorflow-mnist
     stack --docker --docker-image=$IMAGE_NAME build --exec Main
 
+### Docker GPU support
+
+If you want to use GPU you can do:
+
+    IMAGE_NAME=tensorflow/haskell:1.3.0-gpu
+    docker build -t $IMAGE_NAME docker/gpu
+
+We need stack to use nvidia-docker by using a 'docker' wrapper script. This will shadow the normal docker command.
+
+    ln -s `pwd`/tools/nvidia-docker-wrapper.sh <somewhere in your path>/docker
+    stack --docker --docker-image=$IMAGE_NAME setup
+    stack --docker --docker-image=$IMAGE_NAME test
+
 ## Build on macOS
 
 Run the [install_macos_dependencies.sh](./tools/install_macos_dependencies.sh)

docker/Dockerfile

@@ -1,6 +1,6 @@
 # Prepare the image with:
 #   docker build -t tensorflow/haskell:v0 docker
-FROM tensorflow/tensorflow:1.0.0
+FROM tensorflow/tensorflow:1.3.0
 MAINTAINER TensorFlow authors <tensorflow-haskell@googlegroups.com>
 
 RUN apt-get update

docker/gpu/Dockerfile

@@ -0,0 +1,34 @@
+# Prepare the image with:
+#   docker build -t tensorflow/haskell:1.3.0-gpu docker/gpu
+FROM gcr.io/tensorflow/tensorflow:1.3.0-gpu
+MAINTAINER TensorFlow authors <tensorflow-haskell@googlegroups.com>
+
+RUN apt-get update
+
+RUN apt-get install -y \
+        # stack needs git
+        git \
+        # Required by snappy-frames dependency.
+        libsnappy-dev \
+        # Avoids /usr/bin/ld: cannot find -ltinfo
+        libncurses5-dev \
+        # Makes stack viable in the container
+        libgmp-dev \
+        # Required for locales configuration.
+        locales
+
+# Support output of Unicode characters.
+RUN dpkg-reconfigure locales && \
+    locale-gen en_US.UTF-8 && \
+    update-locale LANG=en_US.UTF-8
+
+# Installs protoc and the libraries.
+RUN \
+    curl -O -L https://github.com/google/protobuf/releases/download/v3.2.0/protoc-3.2.0-linux-x86_64.zip && \
+    unzip -d /usr/local protoc-3.2.0-linux-x86_64.zip bin/protoc && \
+    chmod 755 /usr/local/bin/protoc && \
+    curl -O https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-gpu-linux-x86_64-1.3.0.tar.gz && \
+    tar zxf libtensorflow-gpu-linux-x86_64-1.3.0.tar.gz -C /usr/local && \
+    ldconfig
+
+ENV LANG en_US.UTF-8

tools/nvidia-docker-wrapper.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+# This wrapper is necessary to fool 'stack --docker' into
+# using nvidia-docker which again will invoke the real
+# docker in the end.
+#
+# Symlink this file as 'docker' somewhere in path
+
+NV_DOCKER='/usr/bin/docker' nvidia-docker "$@"