Merge pull request Kaggle#815 from Kaggle/add-tf-addons-whl

Build TensorFlow addons wheels.

Author: rosbo

tensorflow-whl/CHANGELOG.md

@@ -14,4 +14,5 @@
 * `2.1.0-py37`: TensorFlow 2.1.0 with Python 3.7
 * `2.1.0-py37-2`: TensorFlow 2.1.0 with Python 3.7 & DLVM base image.
 * `2.1.0-py37-3`: TensorFlow 2.1.0 with Python 3.7, DLVM base image, tensorflow-gcs-config.
-* `2.2.0-py37`: TensorFlow 2.2.0 with Python 3.7.
+* `2.2.0-py37`: TensorFlow 2.2.0 with Python 3.7.
+* `2.2.0-py37-2`: TensorFlow 2.2.0 with Python 3.7 & TFA.

tensorflow-whl/Dockerfile

@@ -62,15 +62,15 @@ RUN apt-get install -y gnupg zip openjdk-8-jdk && \
     dpkg -i bazel_*.deb && \
     rm bazel_*.deb
 
-# Fetch tensorflow & install dependencies.
+# Fetch TensorFlow & install dependencies.
 RUN cd /usr/local/src && \
     git clone https://github.com/tensorflow/tensorflow && \
     cd tensorflow && \
     git checkout tags/v2.2.0 && \
     pip install keras_applications --no-deps && \
     pip install keras_preprocessing --no-deps
 
-# Create a tensorflow wheel for CPU
+# Create a TensorFlow wheel for CPU
 RUN cd /usr/local/src/tensorflow && \
     cat /dev/null | ./configure && \
     bazel build --config=opt \
@@ -80,6 +80,26 @@ RUN cd /usr/local/src/tensorflow && \
     bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_cpu && \
     bazel clean
 
+# Install TensorFlow CPU wheel which is required to build the GCS & addons libraries.
+RUN pip install /tmp/tensorflow_cpu/tensorflow*.whl
+
+# Build TensorFlow addons library against TensorFlow CPU.
+RUN cd /usr/local/src/ && \
+    git clone https://github.com/tensorflow/addons && \
+    cd addons && \
+    git checkout tags/v0.10.0 && \
+    python ./configure.py && \
+    bazel build --enable_runfiles build_pip_pkg && \
+    bazel-bin/build_pip_pkg /tmp/tfa_cpu && \
+    bazel clean
+
+# Build tensorflow_gcs_config library against TensorFlow CPU.
+ADD tensorflow-gcs-config /usr/local/src/tensorflow_gcs_config/
+RUN cd /usr/local/src/tensorflow_gcs_config && \
+    apt-get install -y libcurl4-openssl-dev && \
+    python setup.py bdist_wheel -d /tmp/tensorflow_gcs_config && \
+    bazel clean
+
 # Create a tensorflow wheel for GPU/cuda
 ENV TF_NEED_CUDA=1
 ENV TF_CUDA_VERSION=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION
@@ -103,13 +123,16 @@ RUN cd /usr/local/src/tensorflow && \
     bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_gpu && \
     bazel clean
 
-ADD tensorflow-gcs-config /usr/local/src/tensorflow_gcs_config/
+# Install TensorFlow GPU wheel which to build addons against.
+RUN pip install /tmp/tensorflow_gpu/tensorflow*.whl
 
-# Build tensorflow_gcs_config library against the tensorflow_cpu build
-RUN cd /usr/local/src/tensorflow_gcs_config && \
-    apt-get install -y libcurl4-openssl-dev && \
-    pip install /tmp/tensorflow_cpu/tensorflow*.whl && \
-    python setup.py bdist_wheel -d /tmp/tensorflow_gcs_config && \
+# Build TensorFlow addons library against TensorFlow GPU.
+ENV CUDA_TOOLKIT_PATH=/usr/local/cuda
+ENV CUDNN_INSTALL_PATH=/usr/lib/x86_64-linux-gnu
+RUN cd /usr/local/src/addons && \
+    python ./configure.py && \
+    bazel build --enable_runfiles build_pip_pkg && \
+    bazel-bin/build_pip_pkg /tmp/tfa_gpu && \
     bazel clean
 
 # Print out the built .whl files