Calling Tensorflow models in Scala

In this article I show how to serve a TensorFlow model from Scala¹. I demonstrate this by serializing an example model to the ONNX format, which you can then serve using the Java ONNX runtime² from Scala. Along the way I highlight decisions I made and discuss why you might choose differently for your own work.

Context

ML models need to be served in different ways depending on the functionality they enable. Development and production environments are often separate—this could be to manage security concerns, support different infrastructure (such as phones or IoT devices), or accommodate language or dependency constraints. If the production environment differs from the development environment, you’ll need to explore options for serving the model.

One option is to serialize the model into an intermediary representation format, then serve it using a separate runtime. ONNX is one such format ³. The diagram below shows this process:

Loading diagram…

Exporting to ONNX

I used this timeseries anomaly detection model as an example. Bring your own model if you prefer, but starting with a working example helps you understand the interface. Clone this repo to follow along.

To export a TensorFlow model to ONNX, install tf2onnx. Export the model by running a script from the command line. For my project:

python -m tf2onnx.convert --saved-model ./model --opset 10 --output tfScala/src/main/resources/model.onnx

This produces a single ONNX file. Most arguments are self-explanatory except --opset. The opset number indicates which version of the ONNX API you’re using; the API guarantees that a specific collection of mathematical operations commonly used in machine learning are supported. Now that the model is serialized to ONNX, it’s time to call it from Scala.

Serving the ONNX model in Scala

Runtimes exist in Java and Scala to serve ONNX files. The Scala implementation builds on the Java library and offers type safety and a more accessible abstraction. However, I use the Java library here because it’s more fundamental and makes fewer assumptions about how you’ll interact with models.

Here’s a diagram ⁴ showing how the ONNX runtime integrates within a program to serve the serialized model:

There are some considerations that I’ll call out now:

• ONNX doesn’t currently support operators that manipulate text. Pre-process (tokenize) your textual input before sending it to the model.
• Creating the environment, ONNX tensors, and output all occur as side effects. Decide how to handle them: functional error handling is powerful, but capturing errors at the “ML function didn’t work” level is often sufficient since ONNX’s default errors are usually explicit.
• Cast the ONNX session output to a typed val when retrieving it.

The logic for calling the serialized model through ONNX is in the Interface.scala file. Let’s review each important expression.

Start by loading the model file and instantiating the ONNX environment and session. Here I’ve stored the model file in the jar’s resources. This works if the model is small and stable, but you’d store it externally otherwise.

Create at most one ONNX Environment per Java runtime. You can instantiate multiple sessions (and thus multiple models) as needed. Loading models from bytes is the most reliable approach—it lets you control exactly how the model file is accessed rather than relying on the default createSession implementation. You can also manage hardware and threads through the ONNX session, but that’s beyond this article’s scope.

val modelBytes = Files.readAllBytes(Paths.get(getClass.getResource("model.onnx").toURI))
val env = OrtEnvironment.getEnvironment()
val session = env.createSession(modelBytes)

After creating the environment, convert the input to ONNX tensors. This model accepts an array with dimension , where n is the number of inputs to evaluate. Since this example focuses on integration rather than model content, I’ll create a dummy input of the correct dimension containing zeros.

Note: each value must be wrapped in its own Array, and it must be an Array, not a Seq. This is how the Java ONNX runtime converts values to the fundamentals it requires—see this file.

val input = Array(Seq.fill(288)(0).map(_.toFloat).map(Array(_)).toArray)
val inputTensor = OnnxTensor.createTensor(env, input)
val inputName = "input_1"
val modelInput = Map(inputName -> inputTensor).asJava

Now pass the processed input to the session to run the model. You can unpack the result in several ways, but retrieving output by name makes the code more explicit. Cast the result to a type that’s easier to work with.

val result = session.run(modelInput).get("conv1d_transpose_2")
    .get()
    .getValue
    .asInstanceOf[Array[Array[Array[Float]]]]

Finally, unpack the result for use in the application. For this model, these are anomaly detection probabilities; for an LLM, you’d decode the output into text here. Either way, add error handling and unit tests.

val parsedResult = result.head.flatMap(_.toSeq).toSeq
println(parsedResult)

Next steps

You’ve now seen the workflow for exporting a Python ML model to ONNX format and serving it from Scala. Here are several directions to explore next:

1. Write unit tests for input-to-ONNX-tensor conversions and ONNX-result-to-output parsing. Then add integration tests with a dummy model that matches your real model’s API to verify end-to-end functionality.
2. Functional programming dominates Scala, and justifiably so—its patterns make code easier to maintain. However, the depth of functional abstraction should depend on your business problem and preferences.
3. ONNX achieves significantly higher throughput and lower latency in batch mode.
4. If you need to use this functionality in Spark, rely heavily on the @transient lazy val pattern to manage the model file.

Footnotes

1. This repo contains example code for calling ONNX-serialized models from Scala. ↩

2. If you’re only serving TensorFlow models in Java or Scala, check out TensorFlow Java. I’m tackling the general case—serving any model serializable to ONNX, using a TensorFlow model as an example. TensorFlow Java has a smaller dependency bundle than ONNX; both are much smaller than serving the entire model development environment. ↩

3. Safetensors is another format, but support in the languages I use is limited. Hopefully that will improve. ↩

4. I used Excalidraw to draw this diagram. As mentioned here, I also use Mermaid to create diagrams. ↩