Skip to content

Commit

Permalink
Add support for NVidia GPUs
Browse files Browse the repository at this point in the history
  • Loading branch information
snowzach committed Apr 25, 2020
1 parent 99f5768 commit 2a850c9
Show file tree
Hide file tree
Showing 3 changed files with 154 additions and 0 deletions.
1 change: 1 addition & 0 deletions .dockerignore
Original file line number Diff line number Diff line change
Expand Up @@ -11,4 +11,5 @@ doods
*.zip
private
models/**
Dockerfile.*.*

147 changes: 147 additions & 0 deletions Dockerfile.base.cuda
Original file line number Diff line number Diff line change
@@ -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"]
6 changes: 6 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit 2a850c9

Please sign in to comment.