Loading TensorFlow graphs from Node.js

Check out the related post: Loading a TensorFlow graph with the C++ API.

Even though the full C API for TensorFlow is not yet available, we can still use it load TensorFlow graphs and evaluate them from other languages. This is incredibly useful for embedding pre-trained models in other applications. Embedding is one of the most interesting use cases for TensorFlow as it cannot be accomplished as easily with Theano.

Note that while all of the examples here will use Node.js the steps are nearly identical in any language with C FFI support (e.g. Rust, Go, C#, etc.)

Requirements

• Install Bazel: Google’s build tool used to compile things for TensorFlow.
• Clone the TensorFlow repo.

git clone --recursive https://github.com/tensorflow/tensorflow

Compiling a shared library

We’ll start by compiling a shared library from TensorFlow using Bazel.

UPDATE: The following build rule for creating a shared library is now part of TensorFlow: https://github.com/tensorflow/tensorflow/pull/695

1. Create a new folder in the TensorFlow repo at tensorflow/tensorflow/libtensorflow/.
2. Inside this folder we’re going to create a new BUILD file which will contain a single call to cc_binary with the linkshared option set to 1 so that we get a .so from the build. The name of the binary must end in .so or it will not work.

Here’s the final directory structure:

• tensorflow/tensorflow/libtensorflow/
• tensorflow/tensorflow/libtensorflow/BUILD

Below is the complete BUILD file:

1. From the root of the repository, run ./configure.
2. Compile the shared library with bazel build :libtensorflow.so and locate the generated file from the repo’s root: bazel-bin/tensorflow/libtensorflow/libtensorflow.so

Now that we have our shared library, create a new folder for the host language. Since this is for Node.js I’ll name it tensorflowjs/. This folder can exist outside of the TensorFlow repo since we now have everything needed in the shared library. Copy libtensorflow.so into the new folder.

If you’re on OS X and using Node.js you’ll need to rename the shared library from libtensorflow.so to libtensorflow.dylib. TensorFlow produces an .so however the standard on OS X is dylib. The Node FFI library doesn’t look for .so, only .dylib; however it can read both formats, so we just rename it.

Creating the graph

Just like with the previous C++ tutorial we’re going to create a minimal graph and write it to a protobuf file. (Be sure to name your variables and operations.)

Creating the bindings

Now we can go through the TensorFlow C API header, almost line by line, and write the appropriate binding. Most of the time this is fairly direct, simply copying the signature of the function. I also created variables for many of the common types so they were more legible. For example, any structs which map to void* I declared as variables named after the struct. We can also use the ref-array Node module which provides helpers for types like long long* (essentially an array of long long types) so we’ll define a LongLongArray type to correspond. Otherwise, we just copy the signature:

I also defined a few helper functions to eliminate some of the boilerplate when working with the TensorFlow interface. The first is TF_Destructor, a default tensor destructor for TF_NewTensor. This comment in the TensorFlow source makes it sound like it’s optional but it’s not:

Clients can provide a custom deallocator function so they can pass in memory managed by something like numpy.

Additionally, many TensorFlow functions return a TF_Status struct and checking the status can get tedious. So I defined a function called TF_CheckOK that simply checks if the status code is TF_OK using TF_GetCode. If its not, we throw an error using TF_Message to hopefully get a useful error message. (This function loosely corresponds to TF_CHECK_OK in the TensorFlow source.)

And finally, reading a tensor with TF_TensorData only returns a pointer but to actually read the data we need to extend the returned Buffer to the appropriate length. Creating a Buffer with the correct size is a few lines of boiler plate so I wrapped TF_TensorData to create TF_ReadTensorData which handles that boilerplate for us. Here are the helpers:

Now that we’ve defined our interface the steps for loading the graph are the same as with C++:

1. Initialize a TensorFlow session.
2. Read in the graph we exported above.
3. Add the graph to the session.
4. Setup our inputs and outputs.
5. Run the graph, populating the outputs.
6. Read values from the outputs.
7. Close the session to release resources.

We can load and execute TensorFlow graphs from Node.js! I’ve put the whole thing together into a repo here (you’ll need to provide graph.pb and libtensorflow.dylib since they’re kinda large): https://github.com/jimfleming/tensorflowjs

Follow me on Twitter for more posts like these. We also do applied research to solve machine learning challenges.