Add IoU3D as a custom c++ op (CPU) (#890)

* Initial custom ops testing

* Switch to cpp14

* Add include_package_data flag

* cpp17

* Configure CI tests

* Formatting

* Call configure before running bazel build for CI

* Verbose failures for CI/Bazel

* Switch to tf-cpu for test flow

* Use build dep impls from TFA

* Update cloudbuild workflow to build binaries from source

* Update cloudbuild docks (updated docker image)

* Update GPU tests to TF 2.10

* Update copyright and docstrings

* Move loss to internal, update README and CONTRIBUTING.md

* Update to IoU op ported from Lingvo

* Remove cuda config

* Make iou3d a real keras loss

* Skip custom-ops tests in devcontainer

* Review comments

* Remove skipping of IoU test from devcontainer (it is now opt-in)

* Fix environment variable for windows

* Update environment variable setting

* Use env in yaml for environment variable

* Review comments

* Make IoU a metric, not a loss. Also update Cpp implementation to compute only major axis of pairwise IoU matrix

* typo

* Make test case more demonstrative

* Switch to a single binary, remove binary from git

* Move IoU3D from metrics to ops

* Switch to full matrix computation of iou

* Formatting

.github/CONTRIBUTING.md

@@ -84,6 +84,14 @@ Our default training scripts train using ImageNet. Because we cannot distribute
 
 We look forward to delivering great pre-trained models in KerasCV with the help of your contributions!
 
+## Contributing custom ops
+
+We do not plan to accept contributed custom ops due to the maintenance burden that they introduce. If there is a clear need for a specific custom op that should live in KerasCV, please consult the KerasCV team before implementing it, as we expect to reject contributions of custom ops by default.
+
+We currently support only a small handful of ops that run on CPU and are not used at inference time.
+
+If you are updating existing custom ops, you can re-compile the binaries from source using the instructions in the `Tests that require custom ops` section below.
+
 ## Setup environment
 
 Setting up your KerasCV development environment requires you to fork the KerasCV repository,
@@ -129,6 +137,16 @@ You can run the unit tests for KerasCV by running:
 pytest keras_cv/
 ```
 
+### Tests that require custom ops
+For tests that require custom ops, you'll have to compile the custom ops and make them available to your local Python code:
+```shell
+python build_deps/configure.py
+bazel build keras_cv/custom_ops:all
+cp bazel-bin/keras_cv/custom_ops/*.so keras_cv/custom_ops/
+```
+
+Tests which use custom ops are disabled by default, but can be run by setting the environment variable `TEST_CUSTOM_OPS=true`.
+
 ## Formatting the Code
 We use `flake8`, `isort` and `black` for code formatting.  You can run
 the following commands manually every time you want to format your code:

.github/workflows/actions.yml

@@ -32,9 +32,16 @@ jobs:
           ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
-        pip install tensorflow==2.10.0
+        pip install tensorflow-cpu==2.10.0
         pip install -e ".[tests]" --progress-bar off --upgrade
+    - name: Build custom ops for tests
+      run: |
+        python build_deps/configure.py
+        bazel build keras_cv/custom_ops:all
+        cp bazel-bin/keras_cv/custom_ops/*.so keras_cv/custom_ops/
     - name: Test with pytest
+      env:
+        TEST_CUSTOM_OPS: true
       run: |
         pytest keras_cv/ --ignore keras_cv/models
   format:

BUILD

@@ -0,0 +1,11 @@
+sh_binary(
+    name = "build_pip_pkg",
+    srcs = ["build_deps/build_pip_pkg.sh"],
+    data = [
+        "LICENSE",
+        "MANIFEST.in",
+        "README.md",
+        "setup.py",
+        "//keras_cv",
+    ],
+)

MANIFEST.in

@@ -0,0 +1 @@
+recursive-include keras_cv/custom_ops *.so

README.md

@@ -62,6 +62,23 @@ Thank you to all of our wonderful contributors!
   <img src="https://contrib.rocks/image?repo=keras-team/keras-cv" />
 </a>
 
+
+## Installing Custom Ops from Source
+Installing from source requires the [Bazel](https://bazel.build/) build system
+(version >= 1.0.0).
+
+```
+git clone https://github.com/keras-team/keras-cv.git
+cd keras-cv
+
+python3 build_deps/configure.py
+
+bazel build build_pip_pkg
+bazel-bin/build_pip_pkg wheels
+
+pip install wheels/keras-cv-*.whl
+```
+
 ## Pretrained Weights
 Many models in KerasCV come with pre-trained weights. With the exception of StableDiffusion,
 all of these weights are trained using Keras and KerasCV components and training scripts in this
@@ -72,6 +89,7 @@ history for backbone models [here](examples/training/classification/imagenet/tra
 All results are reproducible using the training scripts in this repository. Pre-trained weights
 operate on images that have been rescaled using a simple `1/255` rescaling layer.
 
+
 ## Citing KerasCV
 
 If KerasCV helps your research, we appreciate your citations.

WORKSPACE

@@ -0,0 +1,3 @@
+load("//build_deps/tf_dependency:tf_configure.bzl", "tf_configure")
+
+tf_configure(name = "local_config_tf")

build_deps/build_pip_pkg.sh

@@ -0,0 +1,88 @@
+#!/usr/bin/env bash
+# Copyright 2022 The KerasCV Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# Builds a wheel of KerasCV for Pip. Requires Bazel.
+# Adapted from https://github.com/tensorflow/addons/blob/master/build_deps/build_pip_pkg.sh
+
+set -e
+set -x
+
+PLATFORM="$(uname -s | tr 'A-Z' 'a-z')"
+function is_windows() {
+  if [[ "${PLATFORM}" =~ (cygwin|mingw32|mingw64|msys)_nt* ]]; then
+    true
+  else
+    false
+  fi
+}
+
+if is_windows; then
+  PIP_FILE_PREFIX="bazel-bin/build_pip_pkg.exe.runfiles/__main__/"
+else
+  PIP_FILE_PREFIX="bazel-bin/build_pip_pkg.runfiles/__main__/"
+fi
+
+function main() {
+  while [[ ! -z "${1}" ]]; do
+    if [[ ${1} == "make" ]]; then
+      echo "Using Makefile to build pip package."
+      PIP_FILE_PREFIX=""
+    else
+      DEST=${1}
+    fi
+    shift
+  done
+
+  if [[ -z ${DEST} ]]; then
+    echo "No destination dir provided"
+    exit 1
+  fi
+
+  # Create the directory, then do dirname on a non-existent file inside it to
+  # give us an absolute paths with tilde characters resolved to the destination
+  # directory.
+  mkdir -p ${DEST}
+  if [[ ${PLATFORM} == "darwin" ]]; then
+    DEST=$(pwd -P)/${DEST}
+  else
+    DEST=$(readlink -f "${DEST}")
+  fi
+  echo "=== destination directory: ${DEST}"
+
+  TMPDIR=$(mktemp -d -t tmp.XXXXXXXXXX)
+
+  echo $(date) : "=== Using tmpdir: ${TMPDIR}"
+
+  echo "=== Copy KerasCV Custom op files"
+
+  cp ${PIP_FILE_PREFIX}setup.py "${TMPDIR}"
+  cp ${PIP_FILE_PREFIX}MANIFEST.in "${TMPDIR}"
+  cp ${PIP_FILE_PREFIX}README.md "${TMPDIR}"
+  cp ${PIP_FILE_PREFIX}LICENSE "${TMPDIR}"
+  rsync -avm -L --exclude='*_test.py' ${PIP_FILE_PREFIX}keras_cv "${TMPDIR}"
+
+  pushd ${TMPDIR}
+  echo $(date) : "=== Building wheel"
+
+  python3 setup.py bdist_wheel > /dev/null
+
+  cp dist/*.whl "${DEST}"
+  popd
+  rm -rf ${TMPDIR}
+  echo $(date) : "=== Output wheel file is in: ${DEST}"
+}
+
+main "$@"

build_deps/configure.py

@@ -0,0 +1,174 @@
+# Copyright 2022 The KerasCV Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Usage: python configure.py
+"""Configures local environment to prepare for building KerasCV from source."""
+
+
+import logging
+import os
+import pathlib
+import platform
+
+import tensorflow as tf
+from packaging.version import Version
+
+_TFA_BAZELRC = ".bazelrc"
+
+
+# Writes variables to bazelrc file
+def write(line):
+    with open(_TFA_BAZELRC, "a") as f:
+        f.write(line + "\n")
+
+
+def write_action_env(var_name, var):
+    write('build --action_env {}="{}"'.format(var_name, var))
+
+
+def is_macos():
+    return platform.system() == "Darwin"
+
+
+def is_windows():
+    return platform.system() == "Windows"
+
+
+def is_linux():
+    return platform.system() == "Linux"
+
+
+def is_raspi_arm():
+    return os.uname()[4] == "armv7l" or os.uname()[4] == "aarch64"
+
+
+def is_linux_ppc64le():
+    return is_linux() and platform.machine() == "ppc64le"
+
+
+def is_linux_x86_64():
+    return is_linux() and platform.machine() == "x86_64"
+
+
+def is_linux_arm():
+    return is_linux() and platform.machine() == "arm"
+
+
+def is_linux_aarch64():
+    return is_linux() and platform.machine() == "aarch64"
+
+
+def is_linux_s390x():
+    return is_linux() and platform.machine() == "s390x"
+
+
+def get_tf_header_dir():
+    import tensorflow as tf
+
+    tf_header_dir = tf.sysconfig.get_compile_flags()[0][2:]
+    if is_windows():
+        tf_header_dir = tf_header_dir.replace("\\", "/")
+    return tf_header_dir
+
+
+def get_cpp_version():
+    cpp_version = "c++14"
+    if Version(tf.__version__) >= Version("2.10"):
+        cpp_version = "c++17"
+    return cpp_version
+
+
+def get_tf_shared_lib_dir():
+    import tensorflow as tf
+
+    # OS Specific parsing
+    if is_windows():
+        tf_shared_lib_dir = tf.sysconfig.get_compile_flags()[0][2:-7] + "python"
+        return tf_shared_lib_dir.replace("\\", "/")
+    elif is_raspi_arm():
+        return tf.sysconfig.get_compile_flags()[0][2:-7] + "python"
+    else:
+        return tf.sysconfig.get_link_flags()[0][2:]
+
+
+# Converts the linkflag namespec to the full shared library name
+def get_shared_lib_name():
+    import tensorflow as tf
+
+    namespec = tf.sysconfig.get_link_flags()
+    if is_macos():
+        # MacOS
+        return "lib" + namespec[1][2:] + ".dylib"
+    elif is_windows():
+        # Windows
+        return "_pywrap_tensorflow_internal.lib"
+    elif is_raspi_arm():
+        # The below command for linux would return an empty list
+        return "_pywrap_tensorflow_internal.so"
+    else:
+        # Linux
+        return namespec[1][3:]
+
+
+def create_build_configuration():
+    print()
+    print("Configuring KerasCV to be built from source...")
+
+    if os.path.isfile(_TFA_BAZELRC):
+        os.remove(_TFA_BAZELRC)
+
+    logging.disable(logging.WARNING)
+
+    write_action_env("TF_HEADER_DIR", get_tf_header_dir())
+    write_action_env("TF_SHARED_LIBRARY_DIR", get_tf_shared_lib_dir())
+    write_action_env("TF_SHARED_LIBRARY_NAME", get_shared_lib_name())
+    write_action_env("TF_CXX11_ABI_FLAG", tf.sysconfig.CXX11_ABI_FLAG)
+
+    # This should be replaced with a call to tf.sysconfig if it's added
+    write_action_env("TF_CPLUSPLUS_VER", get_cpp_version())
+
+    write("build --spawn_strategy=standalone")
+    write("build --strategy=Genrule=standalone")
+    write("build  --experimental_repo_remote_exec")
+    write("build -c opt")
+    write(
+        "build --cxxopt="
+        + '"-D_GLIBCXX_USE_CXX11_ABI="'
+        + str(tf.sysconfig.CXX11_ABI_FLAG)
+    )
+
+    if is_windows():
+        write("build --config=windows")
+        write("build:windows --enable_runfiles")
+        write("build:windows --copt=/experimental:preprocessor")
+        write("build:windows --host_copt=/experimental:preprocessor")
+        write("build:windows --copt=/arch=AVX")
+        write("build:windows --cxxopt=/std:" + get_cpp_version())
+        write("build:windows --host_cxxopt=/std:" + get_cpp_version())
+
+    if is_macos() or is_linux():
+        if not is_linux_ppc64le() and not is_linux_arm() and not is_linux_aarch64():
+            write("build --copt=-mavx")
+        write("build --cxxopt=-std=" + get_cpp_version())
+        write("build --host_cxxopt=-std=" + get_cpp_version())
+
+    print("> Building only CPU ops")
+
+    print()
+    print("Build configurations successfully written to", _TFA_BAZELRC, ":\n")
+    print(pathlib.Path(_TFA_BAZELRC).read_text())
+
+
+if __name__ == "__main__":
+    create_build_configuration()

build_deps/tf_dependency/BUILD.tpl

@@ -0,0 +1,18 @@
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "tf_header_lib",
+    hdrs = [":tf_header_include"],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+)
+
+
+cc_library(
+    name = "libtensorflow_framework",
+    srcs = ["%{TF_SHARED_LIBRARY_NAME}"],
+    visibility = ["//visibility:public"],
+)
+
+%{TF_HEADER_GENRULE}
+%{TF_SHARED_LIBRARY_GENRULE}

build_deps/tf_dependency/build_defs.bzl.tpl

@@ -0,0 +1,4 @@
+# Addons Build Definitions inherited from TensorFlow Core
+
+D_GLIBCXX_USE_CXX11_ABI = "%{tf_cx11_abi}"
+CPLUSPLUS_VERSION = "%{tf_cplusplus_ver}"