Enabling TensorFlow Decision Forests in C

|

TensorFlow Decision Forests is a library that enables TensorFlow to run various decision forest models such as Gradient Boosting Trees (GBT) and Random Forests (RF). In some scenarios (e.g. tabular data) these models are easier to train and comprehend and can outperform deep neural network models (since they have fewer hyper parameters). On the other hand, deep neural networks generally excel in tasks where the data has complex hierarchical structures or involves large amounts of unstructured data, such as images, audio, or text. However, it’s essential to note that the performance of machine learning models can depend on various factors, such as the nature of our data and the complexity of the task.

With that being said, let’s get to it.

If we look at the TF-DF home page, the installation looks pretty straightforward:

!pip install tensorflow tensorflow_decision_forests

But wait, it’s all python. It works well if you want to play around or train models but what if you want to run the trained model in production for inference? I suppose you could create a Python API and interact with your model there but at Samsung Ads on average we are processing 1 million requests per second and have to reply in less than 16 milliseconds. So we needed a high performance system to take care of that requirement and we had two options:

  1. Use the Tensorflow Serving which looks very promising and has minimal overhead.
  2. Load the compiled library in a more performant system (C, Rust, Go, etc.)

Since we already have an existing C/Go system for serving ML models, we decided to go with the second option.

Getting the TF-DF C library

Unfortunately unlike TensorFlow, TF-DF doesn’t come with a straightforward pre-compiled C library (more on this later) it comes with instructions on how to compile it yourself but, it’s not detailed and in order to achieve it, you need to dig into the code yourself and see how everything is tied together. The good news is that TensorFlow has a great Docker image for building TensorFlow components with all the dependencies pre-installed, which is great for getting started. The TF-DF repository contains a script for starting the build container.

Compiling from the source

Let’s start by running the start_compile_docker.sh script and getting the build container running. It downloads the Docker image and gets the build container running. The script also maps the root of the repository to /working_dir of the container. Following the instructions from the compilation guide, it tells us the next step is to run the test_bazel.sh script. However if you run it as is you’ll get the following error message.

ERROR: Could not find a version that satisfies the requirement tensorflow== (from versions: 2.5.0, 2.5.1, 2.5.2, 2.5.3, 2.6.0rc0, 2.6.0rc1, 2.6.0rc2, 2.6.0, 2.6.1, 2.6.2, 2.6.3, 2.6.4, 2.6.5, 2.7.0rc0, 2.7.0rc1, 2.7.0, 2.7.1, 2.7.2, 2.7.3, 2.7.4, 2.8.0rc0, 2.8.0rc1, 2.8.0, 2.8.1, 2.8.2, 2.8.3, 2.8.4, 2.9.0rc0, 2.9.0rc1, 2.9.0rc2, 2.9.0, 2.9.1, 2.9.2, 2.9.3, 2.10.0rc0, 2.10.0rc1, 2.10.0rc2, 2.10.0rc3, 2.10.0, 2.10.1, 2.11.0rc0, 2.11.0rc1, 2.11.0rc2, 2.11.0, 2.11.1, 2.12.0rc0, 2.12.0rc1, 2.12.0, 2.12.1, 2.13.0rc0, 2.13.0rc1, 2.13.0rc2, 2.13.0, 2.13.1, 2.14.0rc0, 2.14.0rc1, 2.14.0, 2.14.1, 2.15.0rc0, 2.15.0rc1, 2.15.0)
ERROR: No matching distribution found for tensorflow==


Oops, looks like we need to specify the version but it’s not explained in the guide. Upon inspection of the script, we can see that it requires a few variables and it’s nicely documented at the top of the script:

# Build and test TF-DF.
# Options
#  RUN_TESTS: Run the unit tests e.g. 0 or 1.
#  PY_VERSION: Version of Python to be used, must be at least 3.9
#  STARTUP_FLAGS: Any flags given to bazel on startup
#  TF_VERSION: Tensorflow version to use or "nightly".
#              For cross-compiling with Apple Silicon for Mac Intel, use 
#              mac-intel-crosscompile.
#              Tests will not work when cross-compiling (obviously).
# FULL_COMPILATION: If 1, compile all parts of TF-DF. This may take a long time.
#
# Usage example
#
#   RUN_TESTS=1 PY_VERSION=3.9 TF_VERSION=2.13.0 ./tools/test_bazel.sh


Let’s run it with the suggested variables, just make sure to run it at the root of the project, otherwise it returns an error saying that it can’t find the WORKSPACE file:

$ RUN_TESTS=1 PY_VERSION=3.9 TF_VERSION=2.13.0 ./tools/test_bazel.sh


On an 8 core machine, it takes about 30 minutes to compile. After the compilation is done, we can find the compiled library at tensorflow/ops/inference/inference.so. We copy the compiled library to /usr/local/lib by convention.

Getting it from the Python package

For us, compiling TF-DF was not the best idea. It takes a very long time to compile and it’s very error prone. An easier way to get it is through the official TF-DF pypi package. These wheel packages are essentially zip archives. You can download the specific version for the platform that you want, extract it and find the compiled library inside.

$ wget  https://files.pythonhosted.org/packages/d4/80/d60d714fede4c53542b75bf731da9328350b00c77bc58e9ac882697055a9/tensorflow_decision_forests-1.8.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
$ tar -xzf tensorflow_decision_forests-1.8.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl


If we dig into the contents of this archive, this is what we see:

├── __init__.py
├── component
│  ├── __init__.py
│  ├── builder
│  │  ├── __init__.py
│  │  ├── builder.py
│  │  └── builder_test.py
│  ├── inspector
│  │  ├── __init__.py
│  │  ├── blob_sequence.py
│  │  ├── blob_sequence_test.py
│  │  ├── inspector.py
│  │  └── inspector_test.py
│  ├── model_plotter
│  │  ├── __init__.py
│  │  ├── model_plotter.py
│  │  ├── model_plotter_test.py
│  │  └── plotter.js
│  ├── py_tree
│  │  ├── __init__.py
│  │  ├── condition.py
│  │  ├── condition_test.py
│  │  ├── dataspec.py
│  │  ├── dataspec_test.py
│  │  ├── node.py
│  │  ├── node_test.py
│  │  ├── objective.py
│  │  ├── objective_test.py
│  │  ├── tree.py
│  │  ├── tree_test.py
│  │  ├── value.py
│  │  └── value_test.py
│  └── tuner
│ 	├── __init__.py
│ 	├── tuner.py
│ 	└── tuner_test.py
├── keras
│  ├── __init__.py
│  ├── core.py
│  ├── core_inference.py
│  ├── grpc_worker_main
│  ├── keras_distributed_test.py
│  ├── keras_test.py
│  ├── keras_tuner_test.py
│  ├── test_runner.py
│  ├── wrappers.py
│  └── wrappers_pre_generated.py
└── tensorflow
   ├── __init__.py
   ├── cc_logging.py
   ├── check_version.py
   ├── check_version_test.py
   ├── core.py
   ├── core_inference.py
   ├── core_test.py
   ├── ops
   │  ├── __init__.py
   │  ├── inference
   │  │  ├── __init__.py
   │  │  ├── api.py
   │  │  ├── inference.so
   │  │  ├── op.py
   │  │  ├── op_dynamic.py
   │  │  ├── test_utils.py
   │  │  ├── tf1_test.py
   │  │  └── tf2_test.py
   │  └── training
   │ 	├── __init__.py
   │ 	├── op.py
   │ 	├── op_dynamic.py
   │ 	├── op_test.py
   │ 	└── training.so
   ├── tf1_compatibility.py
   └── tf_logging.py


As you can see the compiled library can be found at tensorflow/ops/inference/inference.so.

Running it

Now that we have the compiled library, how do we load it? Unfortunately we couldn’t find any documentation around this subject. What we learned is that, TF-DF hooks into TensorFlow using custom operators and unlike TensorFlow itself, TF-DF is not automatically picked up. It can be loaded dynamically by calling the TF_LoadLibrary function.
Here’s a very simple C program that loads TF-DF if it’s available:

#include <stdio.h>
#include <unistd.h>
#include <tensorflow/c/c_api.h>

void load_tfdf()
{
    char *lib_path = "/usr/local/lib/inference.so";

    // check if file exists
    if (access(lib_path, F_OK) != 0)
    {
        printf("Can't find TFDF lib at %s\n", lib_path);
        return;
    }

    printf("Loading TFDF...");

    TF_Status *status = TF_NewStatus();
    TF_LoadLibrary(lib_path, status);
    TF_Code code = TF_GetCode(status);
    if (code == TF_OK)
    {
        printf("OK!\n");
    }
    else
    {
        printf("ERROR.\n");
    }
}

int main()
{
    printf("Hello from Tensorflow C library version %s\n", TF_Version());
    load_tfdf();

    return 0;
}


After compiling and running it, we get the following output:

$ gcc -I/usr/local/include -L/usr/local/lib hello_tf.c -ltensorflow -o hello_tf && ./hello_tf

Hello from Tensorflow C library version 2.13.0
Loading TFDF...OK!

Nice! Looks like TF-DF is loaded and now we can load our decision forest models.

TF-DF is a great extension to Tensorflow. It enabled us to add support for decision forest models within our existing code base with minimal changes and the team is very happy with how well these models are performing compared to DNNs. Also we didn’t need to introduce a new dependency and additional overhead such as TensorFlow Serving in order to be able to serve these models. Because of the lack of documentation, integrating it with our existing C code base however was quite challenging and was not as simple as the Tensorflow itself. We hope that this post could help with filling that gap and provide a bit of guidance to other developers facing similar challenges.