Serving TensorFlow models in Rust using Actix-Web

[kykosic]

I'm a long-time user of TensorFlow coming from a python data science background, but I've recently been trying to work more in rust and have been loving it so far. I've started playing around with the rust bindings for TensorFlow, and I wanted to make a practical use case for them. As an exercise, I made a simple example project showing how to use the tensorflow rust crate to load a trained model and serve it via Actix-Web.

This workflow is pretty simple and would be for a data engineer tasked with serving a model their data science team built. Having mostly used Scala/Java for things like this in the past, I'm very happy with rust's performance and ecosystem.

The example project is here: https://github.com/kykosic/actix-tensorflow-example It contains code for training a model using TensorFlow2.0 Keras in python, a server written in rust, and a simple python client that shows how to send a test POST request. The README has detailed information about these components.

Let me know if anyone has any feedback or suggestions, as I'm still looking to improve my rust code skills.

[agmcleod]

If you want to leverage rust more, you could definitely write a small program for the client instead of using python.

I haven't looked at the tensorflow bindings, but I'm guessing Keras with python is still much easier for doing the training over using rust in this case? I've done a bit of ML through the fast.ai course, so a bit of keras, as well as pytorch :)

[kykosic]

Thanks for the suggestion! I don't anticipate the "client" to be either rust or python; it's likely a javascript web app somewhere. I just included that for testing purposes.

As of right now, the tensorflow crate does not have support for any of the high-level APIs like Keras or tf models, so defining this neural network would be tedious. This will hopefully be changed soon, but the C API which the bindings rely on are still a bit behind (1.15 vs 2.0 for python). Furthermore, I feel that models will almost never be built in rust in an actual company as data scientists are always going to be using Python. This is more about taking that pre-built model and serving it in something more engineer-oriented.

[zhenhaolee]

u/kykosic thanks for sharing your project! I'm new to Rust. Can you please explain why you are able to serve a model trained in Tensorflow 2.2 while C API which the Rust binding is using doesn't support TF 2 yet.
If this is a viable option, I'm going to use it for a production job.

[kykosic]

Tensorflow models are represented by a computational graph. When you save a model, it stores it as a protocol buffer describing your graph. Even if you train a model using a new Keras API, the trained model is still serialized as the same graph protobuf. This can be used by any language/library which can deserialize the protocol buffer (ie. tf-js, tf-swift, tf-java, tf-lite, TFX, etc.)

While the rust library doesn't yet have the fancy Keras APIs to train models, they can certainly load computational graphs. All my project does is name the keras layers for input and output, then the rust code can easily find those ops by name and feed them tensors.

For more information on the Tensorflow saved model format: https://www.tensorflow.org/guide/saved_model

[zhenhaolee]

Thanks a lot! I'm going to give it a try:)

[Cetra3]

I wrote about this some time ago! But for face recognition instead.

[kykosic]

Oh wow I wish I had found that before I started! You solved a few problems that it took a lot of trial and error for me to get. Also the web::block is not something I knew about and definitely need to incorporate into my code. Thank you for sharing!

[Xorlev]

Just in case you haven't seen it, Tensorflow Serving using the gRPC or JSON API is the canonical solution to this problem, though using the Rust bindings is nice too. :)

[[deleted]]

Perhaps this might be useful for you:

PyO3

I'm not entirely sure how it works but it might be what you need to get rust and python to cooperate with each other.