How do TensorFlow.js models load in the browser?

Understanding the process of loading TensorFlow.js models in web browsers with TensorFlow Hub

In a previous post I wrote about TensorFlow.js and how it can be used to build and deploy deep learning models in the browser. In this post I will focus on how TensorFlow.js models are loaded in the browser.

Most of the times Machine Learning models are used from an API. But TensorFlow.js brings those models to the user browser. The process behind this is very interesting.

Toxicity Model

There are many ways to load and use TensorFlow.js models.

Let's say you want to load the Toxicity model. Toxicity is a model that can detect toxicity in text.

This is the kind of model you can load directly from an NPM package after install it.

import * as toxicity from "@tensorflow-models/toxicity";

So, now you can load the model by calling toxicity.load().

Toxicity.load() it's a wrapper around tf.loadGraphModel().

tf.loadGraphModel() is a function that loads a TensorFlow.js model from a URL or a local file.

Process of loading the model

Even when you install a model with NPM, the model is not inside the package. When you need to use the model this will fire a network call to download the model.

Because the model is loaded asynchronously, you need to wait for the model to load before you can use it.

const model = await toxicity.load();

This is something we can see from the Network tab in Chrome Developer Tools:

Shards

The Toxicity model is a large model, so it's split into multiple files. The files are called shards. The shards are loaded in parallel, and then they are combined into a single model.

This allows you to load the model more efficiently, especially when loading large models over a slow connection. TensorFlow.js automatically handles the process of downloading and combining the shards to create the full model.

Not all the models are split into shards, some models can be loaded as single files because they are not very large.

Storage

This model will be storaged in local memory, so you don't need to worry about downloading the model each time you run your application.

Once the model is complete, you can start using it through your app.

This was the internal process of loading machine learning models in the browser using TensorFlow.js.