AutoML Vision Edge: Loading and Running a TensorFlow.js Model (Part 2)

Running TensorFlow.js models on a Node.js server

Photo by Kahika on Unsplash

This post is the 4th post in my series on Google Cloud’s AutoML Vision Edge. The post will cover using the Tensorflow.js models on NodeJS servers. If you are not following the AutoML Vision Edge series till now, here is some recap:

We started the series by learning how we can train and run inference on an edge-ready ML model in hours using Google Cloud’s AutoML.

In the first post, we used the .tflite format provided by AutoML to run inference on-device.

AutoML Vision Edge: Build machine learning models for mobile and edge devices-in hours

In the second post, we used Python to run inference on the TF Saved Model format, available to us as .pb file.

AutoML Vision Edge: Exporting and Loading TensorFlow SavedModels with Python

In the third post, we used the tf.js format and to load and run the model directly in the browser.

AutoML Vision Edge: Loading and Running a TensorFlow.js Model (Part 1)

In this post, we will learn how to use the Tf.js format of the AutoML model on the server instead of in the browser. If you’re new to this series, I’d recommend going through part 1 of the series before reading ahead.

TensorFlow.js with Node.js

TensorFlow.js was extended to run on Node.js, using an extension library called tfjs-node. This allows TensorFlow.js to be used in backend JavaScript applications without having to use Python.

If you’re working in web development or working on a mobile application that uses Node.js servers, you can easily use AutoML with tfjs-node to train and run inference within your application. This will remove any additional dependencies like changing the language to Python (or looking for an additional Python developer to complete the project). The client and server both can be written in JavaScript and the ML models can also be inferred from the server-side.

Node.js support within with AutoML is not properly documented anywhere. This post will help you to understand and use the tfjs-node package with AutoML-based edge models.

You can learn more about this here.

TensorFlow.js & Node.js with AutoML

To use AutoML models in Node.js, the first step is to install all the packages required.

Step 1: Installing Packages

npm init
npm install @tensorflow/tfjs-node @tensorflow/tfjs-automl 

Step 2: Including all packages at the top of the file:

const tf = require("@tensorflow/tfjs-node");
const automl = require("@tensorflow/tfjs-automl");
const fs = require("fs");
const path = require('path');

In a TensorFlow.js model, as we have seen in the previous blog post, there are many .bin files as well as a model.json file.

Step 3: Loading the model and the dictionary file

tf.js also requires a dictionary file in which we’ve defined the classes for the classification model. The dictionary file is available with the exported tf.js model format:

const loadDictionary = modelUrl => {
   const lastIndexOfSlash = modelUrl.lastIndexOf("/");
   const prefixUrl = lastIndexOfSlash >= 0 ? modelUrl.slice(0,
   lastIndexOfSlash + 1) : "";
   const dictUrl = `${prefixUrl}dict.txt`;
   const text = fs.readFileSync(dictUrl, { encoding: "utf-8" });
   return text.trim().split("\n");
 };

We’ll use tf.loadGraphModel() from tensorflow/tfjs-node. The function will load all the model files or the .bin files using the model.json files.

The function returns a Promise<tf.GraphModel>, which is a directed, acyclic graph built from the SavedModel GraphDef and allows us to execute inference:

const [model, dict] = await     Promise.all([tf.loadGraphModel(`file://${modelUrl}`),
loadDictionary(modelUrl)
]);

Once the Promise chain provides us with the model and dict variables, we can then pass them to the AutoML ImageClassificationModel function, as shown below:

const loadImageClassification = async modelUrl => {
  const [model, dict] = await Promise.all([
                       tf.loadGraphModel(`file://${modelUrl}`),
                       loadDictionary(modelUrl)
  ]);
  return new automl.ImageClassificationModel(model, dict);
};

Step 4: Running inference on the input image

Before running inference, we need to make sure we’re providing the right format to the model. In this case, our model accepts the 3d tensor of a Uint8Array. You can use Netron to find out the input format supported by our model.

const decodeImage = imgPath => {
  const imgSrc = fs.readFileSync(imgPath);
  const arrByte = Uint8Array.from(Buffer.from(imgSrc));
  return tf.node.decodeImage(arrByte);
};
const main = async()=>{
  await loadImageClassification(modelURL);
  return await model.classify(decodedImage);
}

Finally, we’ll write our main function, which will call our loadModelClassification function. And then we’ll run model.classify, which is the inference method on the decoded image.

Now, we can call the entire script as shown below:

let output = await main()
                   .then((val)=>{
                         return val
                    }).catch(console.log);
console.log(output);

The output will be a JavaScript object containing the labels from the dictionary file with their respective probabilities:

{"daisy": 0.12, "lotus": 99.88}

Conclusion

In this post, we’ve seen how we can use the TensorFlow.js edge model format provided by AutoML on Node.js servers. Though there’s very little documentation of this on the internet, following the above steps gives impressive results on the server as well.

Article 3 and article 4 [this post] of this series remove any restrictions for the TF.js format to be used on the client or the server-side.

In upcoming posts of this series, we’ll learn about other formats supported by AutoML and their respective use-cases.

If you liked the article, please clap your heart out. Tip — Your 50 claps will make my day!

Want to know more about me? Please check out my website. If you’d like to get updates, follow me on Twitter and Medium. If anything isn’t clear or you want to point out something, please comment down below.

Editor’s Note: Heartbeat is a contributor-driven online publication and community dedicated to providing premier educational resources for data science, machine learning, and deep learning practitioners. We’re committed to supporting and inspiring developers and engineers from all walks of life.

Editorially independent, Heartbeat is sponsored and published by Comet, an MLOps platform that enables data scientists & ML teams to track, compare, explain, & optimize their experiments. We pay our contributors, and we don’t sell ads.

If you’d like to contribute, head on over to our call for contributors. You can also sign up to receive our weekly newsletters (Deep Learning Weekly and the Comet Newsletter), join us on Slack, and follow Comet on Twitter and LinkedIn for resources, events, and much more that will help you build better ML models, faster.