From 2a850c99740c30b19ba7d653751d5174f0a0253c Mon Sep 17 00:00:00 2001
From: Zach Brown <git@prozach.org>
Date: Sat, 25 Apr 2020 05:44:14 +0000
Subject: [PATCH] Add support for NVidia GPUs

---
 .dockerignore        |   1 +
 Dockerfile.base.cuda | 147 +++++++++++++++++++++++++++++++++++++++++++
 README.md            |   6 ++
 3 files changed, 154 insertions(+)
 create mode 100644 Dockerfile.base.cuda

diff --git a/.dockerignore b/.dockerignore
index bc8fd46..ea2c18a 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -11,4 +11,5 @@ doods
 *.zip
 private
 models/**
+Dockerfile.*.*
 
diff --git a/Dockerfile.base.cuda b/Dockerfile.base.cuda
new file mode 100644
index 0000000..d84f924
--- /dev/null
+++ b/Dockerfile.base.cuda
@@ -0,0 +1,147 @@
+FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04 as builder
+
+# Install reqs with cross compile support
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    pkg-config zip zlib1g-dev unzip wget bash-completion git curl \
+    build-essential patch g++ python python-future python-numpy python-six python3 \
+    cmake ca-certificates \
+    libc6-dev libstdc++6 libusb-1.0-0
+
+# Install protoc
+RUN wget https://github.com/protocolbuffers/protobuf/releases/download/v3.9.1/protoc-3.9.1-linux-x86_64.zip && \
+    unzip protoc-3.9.1-linux-x86_64.zip -d /usr/local && \
+    rm /usr/local/readme.txt && \
+    rm protoc-3.9.1-linux-x86_64.zip
+
+# Install bazel
+RUN wget https://github.com/bazelbuild/bazel/releases/download/0.27.1/bazel_0.27.1-linux-x86_64.deb && \
+    dpkg -i bazel_0.27.1-linux-x86_64.deb && \
+    rm bazel_0.27.1-linux-x86_64.deb
+
+# Download tensorflow sources
+ARG TF_VERSION="v2.1.0"
+ENV TF_VERSION $TF_VERSION
+RUN cd /opt && git clone https://github.com/tensorflow/tensorflow.git --branch $TF_VERSION --single-branch
+
+# Configure tensorflow
+ENV TF_NEED_GDR=0 TF_NEED_AWS=0 TF_NEED_GCP=0 TF_NEED_CUDA=1 TF_NEED_HDFS=0 TF_NEED_OPENCL_SYCL=0 TF_NEED_VERBS=0 TF_NEED_MPI=0 TF_NEED_MKL=0 TF_NEED_JEMALLOC=1 TF_ENABLE_XLA=0 TF_NEED_S3=0 TF_NEED_KAFKA=0 TF_NEED_IGNITE=0 TF_NEED_ROCM=0
+RUN cd /opt/tensorflow && yes '' | ./configure
+
+# Tensorflow build flags for rpi
+ENV BAZEL_COPT_FLAGS="--local_resources 16000,16,1 --config monolithic --copt=-O3 --copt=-fomit-frame-pointer --copt=-mfpmath=both --copt=-mavx --copt=-msse4.2 --incompatible_no_support_tools_in_action_inputs=false --config=noaws --config=nohdfs"
+ENV BAZEL_EXTRA_FLAGS=""
+
+# Patch to make it work with cuda 10.2
+RUN cd /opt/tensorflow && sed -i '/"--bin2c-path=%s" % bin2c.dirname,/d' third_party/nccl/build_defs.bzl.tpl
+
+# Compile and build tensorflow lite
+RUN cd /opt/tensorflow && \
+    bazel build -c opt $BAZEL_COPT_FLAGS --verbose_failures $BAZEL_EXTRA_FLAGS //tensorflow/lite:libtensorflowlite.so && \
+    install bazel-bin/tensorflow/lite/libtensorflowlite.so /usr/local/lib/libtensorflowlite.so && \
+    bazel build -c opt $BAZEL_COPT_FLAGS --verbose_failures $BAZEL_EXTRA_FLAGS //tensorflow/lite/experimental/c:libtensorflowlite_c.so && \
+    install bazel-bin/tensorflow/lite/experimental/c/libtensorflowlite_c.so /usr/local/lib/libtensorflowlite_c.so && \
+    mkdir -p /usr/local/include/flatbuffers && cp bazel-tensorflow/external/flatbuffers/include/flatbuffers/* /usr/local/include/flatbuffers
+
+# Compile and install tensorflow shared library
+RUN cd /opt/tensorflow && \
+    bazel build -c opt $BAZEL_COPT_FLAGS --verbose_failures $BAZEL_EXTRA_FLAGS //tensorflow:libtensorflow.so && \
+    install bazel-bin/tensorflow/libtensorflow.so /usr/local/lib/libtensorflow.so && \
+    ln -rs /usr/local/lib/libtensorflow.so /usr/local/lib/libtensorflow.so.1
+
+# cleanup so the cache directory isn't huge
+RUN cd /opt/tensorflow && \
+    bazel clean && rm -Rf /root/.cache
+
+# Install GOCV
+ARG OPENCV_VERSION="4.3.0"
+ENV OPENCV_VERSION $OPENCV_VERSION
+RUN cd /tmp && \
+    curl -Lo opencv.zip https://github.com/opencv/opencv/archive/${OPENCV_VERSION}.zip && \
+    unzip -q opencv.zip && \
+    curl -Lo opencv_contrib.zip https://github.com/opencv/opencv_contrib/archive/${OPENCV_VERSION}.zip && \
+    unzip -q opencv_contrib.zip && \
+    rm opencv.zip opencv_contrib.zip && \
+    cd opencv-${OPENCV_VERSION} && \
+    mkdir build && cd build && \
+    cmake -D CMAKE_BUILD_TYPE=RELEASE \
+    -D CMAKE_INSTALL_PREFIX=/usr/local \
+    -D OPENCV_EXTRA_MODULES_PATH=../../opencv_contrib-${OPENCV_VERSION}/modules \
+    -D WITH_JASPER=OFF \
+    -D WITH_QT=OFF \
+    -D WITH_GTK=OFF \
+    -D BUILD_DOCS=OFF \
+    -D BUILD_EXAMPLES=OFF \
+    -D BUILD_TESTS=OFF \
+    -D BUILD_PERF_TESTS=OFF \
+    -D BUILD_opencv_java=NO \
+    -D BUILD_opencv_python=NO \
+    -D BUILD_opencv_python2=NO \
+    -D BUILD_opencv_python3=NO \
+    -D OPENCV_GENERATE_PKGCONFIG=ON .. && \
+    make -j $(nproc --all) && \
+    make preinstall && make install && \
+    cd /tmp && rm -rf opencv*
+
+# Download the edgetpu library and install it
+RUN cd /tmp && git clone https://github.com/google-coral/edgetpu.git && \
+    install edgetpu/libedgetpu/throttled/k8/libedgetpu.so.1.0 /usr/local/lib/libedgetpu.so.1.0 && \
+    ln -rs /usr/local/lib/libedgetpu.so.1.0 /usr/local/lib/libedgetpu.so.1 && \
+    ln -rs /usr/local/lib/libedgetpu.so.1.0 /usr/local/lib/libedgetpu.so && \
+    mkdir -p /usr/local/include/libedgetpu && \
+    install edgetpu/libedgetpu/edgetpu.h /usr/local/include/libedgetpu/edgetpu.h && \
+    install edgetpu/libedgetpu/edgetpu_c.h /usr/local/include/libedgetpu/edgetpu_c.h && \
+    rm -Rf edgetpu
+
+# Configure the Go version to be used
+ENV GO_ARCH "amd64"
+ENV GOARCH=amd64
+
+# Install Go
+ENV GO_VERSION "1.14.2"
+RUN curl -kLo go${GO_VERSION}.linux-${GO_ARCH}.tar.gz https://dl.google.com/go/go${GO_VERSION}.linux-${GO_ARCH}.tar.gz && \
+    tar -C /usr/local -xzf go${GO_VERSION}.linux-${GO_ARCH}.tar.gz && \
+    rm go${GO_VERSION}.linux-${GO_ARCH}.tar.gz
+
+FROM nvidia/cuda:10.2-cudnn7-runtime-ubuntu18.04 as build
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    pkg-config zip zlib1g-dev unzip wget bash-completion git curl \
+    build-essential patch g++ python python-future python3 ca-certificates \
+    libc6-dev libstdc++6 libusb-1.0-0
+
+# Copy all libraries, includes and go
+COPY --from=builder /usr/local/. /usr/local/.
+COPY --from=builder /opt/tensorflow /opt/tensorflow 
+
+ENV GOOS=linux
+ENV CGO_ENABLED=1
+ENV CGO_CFLAGS=-I/opt/tensorflow
+ENV PATH /usr/local/go/bin:/go/bin:${PATH}
+ENV GOPATH /go
+
+# Create the build directory
+RUN mkdir /build
+WORKDIR /build
+
+ADD . .
+RUN make
+
+FROM nvidia/cuda:10.2-cudnn7-runtime-ubuntu18.04
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends libusb-1.0 libc++-7-dev wget unzip ca-certificates libdc1394-22 libavcodec57 libavformat57 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+RUN mkdir -p /opt/doods
+WORKDIR /opt/doods
+COPY --from=build /usr/local/lib/. /usr/local/lib/.
+COPY --from=build /build/doods /opt/doods/doods
+ADD config.yaml /opt/doods/config.yaml
+ENV LD_LIBRARY_PATH=/usr/local/cuda-10.2/compat
+RUN ldconfig
+
+RUN mkdir models
+RUN wget https://storage.googleapis.com/download.tensorflow.org/models/tflite/coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip && unzip coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip && rm coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip && mv detect.tflite models/coco_ssd_mobilenet_v1_1.0_quant.tflite && rm labelmap.txt
+RUN wget https://dl.google.com/coral/canned_models/coco_labels.txt && mv coco_labels.txt models/coco_labels0.txt
+
+CMD ["/opt/doods/doods", "-c", "/opt/doods/config.yaml", "api"]
diff --git a/README.md b/README.md
index 910c6a3..8b7be58 100644
--- a/README.md
+++ b/README.md
@@ -171,6 +171,12 @@ There are several published Docker images that you can use
 * aarch64 - Arm 64 bit image
 * rpi - Arm 32 bit/arm7 image optimized for the Raspberry Pi
 * amd64 - 64 bit x86 image with all the fancy cpu features like avx and sse4.2
+* cuda - Support for NVidia GPUs
+
+## CUDA Support
+There is now NVidia GPU support with an docker image tagged cuda, to run:
+`docker run -it --gpus all -p 8080:8080 snowzach/doods:cuda`
+For whatever reason, it can take a good 60-80 seconds before the model finishes loading.
 
 ## Compiling
 You can compile it yourself using the native config which should optimize for whatever you have installed.