Serving TensorFlow model in Scala
By: Kirill Panarin
Tags: Machine Learning, Data Science, Scala, Functional Programming, TensorFlow

Taken from here

Taken from here

Problem

The process of training a machine learning model is only the first step towards putting it into production. At the end of the day, data science team sign-offs a model and data engineering team has to figure out how to deploy, use, update and maintain the model in the production environment. There are two possible approaches on how to do this:

Expose a model through the RPC/REST API. The advantage, in this case, is that a language/infrastructure used to build the model could be also used to serve RPC/REST API. The downside is that every request to a model is a separate RPC or HTTP call with its own latency. The latency puts an upper limit on an average time it takes to query the model.

Use the model directly in the data engineering pipeline already available in the production environment. In this case, there is no issue with an extra latency, but, it appears, the best data processing engines/frameworks do not support the same set of languages used for models training. For instance, the typical scenario is when a data science team builds a model using Python/TensorFlow, and the data engineering team has to integrate this into Spark/Scala/Java stack.

I am going to focus on the second option and describe how to build a library for using a pre-trained TensorFlow model in Scala. TensorFlow already provides Java API, which could be directly referenced and used in Scala, but it is very low level and cumbersome.

The resulting library with examples described here is available in tfModelServing4s GitHub repo.

API Algebra

Tagless Final approach is used to build the library. More details about it could be found here and here. The idea is to create an algebra which abstracts over an effect and types used in a concrete implementation/interpreter. This trait defines the set of operations the API is exposing:

F[_] is the high-kinded type representing an effect a concrete implementation is using. This could be a Future, Scalaz Task or even simple Try. The trait also defines two types TModel and TTensor. The former is a model type that is going to represent a loaded into the memory model and the latter a tensor type. Both types have to be defined by an interpreter. In addition, the trait defines several basic operations for working with a model:

load is for loading a model from some source e.g. from the file system. ModelSource is a co-product type defining the model source:

metadata queries model metadata from the loaded model. It returns ModelMetadata which, in its turn, enumerates list of signatures available for this graph. The notion of signature is borrowed from TensorFlow and it defines a function exposed by a model. For instance, a prediction model might expose a signature which is defined by its input and output tensors. When an input tensor is given the function produces an output tensor which might by a probability of an input to belong to a particular class. A graph might expose multiple signatures if it supports more than one task. Here is the definition of ModelMetadata, SignatureMetadata and TensorMetadata case classes:

tensor function creates a tensor object from a compatible data structure e.g. Array or List. The data structure is compatible if it has an implicit TensorEncoder type class defined for it:

Here is an example of the TensorEncoder for converting 2-dimensional Array to TensorFlow API Tensor object:

eval is the function which actually executes a graph computation. It expects a model object, a metadata for the output and the map with input tensor metadata as a key and an input tensor itself as a value. It returns the resulting tensor object wrapped into an effect F[_]. The function also requires TensorDecoder and Closeable type classes for TTensor type. TensorDecoder converts a tensor object into its representation in a form of a data structure like Array, List etc.:

Closeable type class is needed in order to dispose an unmanaged object. For example, TensorFlow API operates on tensors which are un-managed objects and they have to be explicitly closed when they are not needed anymore, otherwise, they going to stay in the memory and this might lead to memory leaks in the application:

That is it. The API is defined. The next step is to implement an “interpreter” for TensorFlow.

TensorFlow Interpreter

The interpreter for TensorFlow has to implement ModelServing trait. I have chosen scala.util.Try as an effect so that all API functions will return results wrapped into Try and it will be possible to compose API operations using Scala for comprehension. Here is the implementation of the interpreter. It does all dirty work of interacting with TensorFlow Java API:

In addition, we have to define TensorEncoder, TensorDecoder and Closeable type classes:

There are encoders and decoders only for a string, 1-dimensional and 2-dimensional Array. This is the basic minimum I needed to execute several examples and see if it works. This could be extended for any other compatible data type.

Usage Example #1

It is time to try the library on a simple example. First, a TensorFlow saved model has to be prepared. For this example, I have used a computational graph which simply multiplies a matrix with the identity matrix. The output matrix should be equal to the input matrix according to the identity matrix property. Here is the Python code building the graph and saving it in TensorFlow SavedModel format:

It defines a signature with the single input (‘x’) and the single output (‘y’). The resulting model directory has the serialized graph together with the variables snapshot:

The Scala code feeding an input matrix into the model and getting the result matrix:

Here is the output produced by the application above:

Usage Example #2

The second usage example is based on the dog breed classification model I described how to build here. It takes a dog image and produces top 5 breed probabilities:

The output produced for a sample dog image:

Conclusion

A reasonable Scala library to work with TensorFlow models was built. It could be used to perform classifications, predictions etc.using a pre-trained TensorFlow model in any Scala/Java application e.g. Apache Spark.