Deploy Tensoflow model on browser

For learning purpose, I deployed my trained model using Tensorflow trained on MNIST data set and deploy it on browser. The browser then can preprocess input image and make predictions without any extra calls to server. In practice this can be used with computer vision for lowering latancy. But we don’t really want to make our model publicly available to anyone.

This is how it looks like.

First, we build a simple model. Import some common libraries. We use tfjs to export model to json at the final step.

import os

import tensorflow as tf
import tensorflowjs as tfjs
from tensorflow import keras

We use MNIST data set for this example.

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

train_labels = train_labels[:1000]
test_labels = test_labels[:1000]

train_images = train_images[:1000].reshape(-1, 28 * 28) / 255.0
test_images = test_images[:1000].reshape(-1, 28 * 28) / 255.0

Create a simple 3 layers NN model.

# Define a simple sequential model
def create_model():
  model = tf.keras.Sequential([
    keras.layers.Dense(512, activation='relu', input_shape=(784,)),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(10, activation='sigmoid'),
  ])

  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

  return model

# Create a basic model instance
model = create_model()

# Display the model's architecture
model.summary()

Model: "sequential"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 dense (Dense)               (None, 512)               401920    
                                                                 
 dropout (Dropout)           (None, 512)               0         
                                                                 
 dense_1 (Dense)             (None, 10)                5130      
                                                                 
=================================================================
Total params: 407,050
Trainable params: 407,050
Non-trainable params: 0
____________________________

model.fit(train_images, train_labels, epochs=5)
tfjs.converters.save_keras_model(model, 'tfmnist_json')

Epoch 1/5
32/32 [==============================] - 0s 3ms/step - loss: 1.1469 - sparse_categorical_accuracy: 0.6740
Epoch 2/5
32/32 [==============================] - 0s 3ms/step - loss: 0.4375 - sparse_categorical_accuracy: 0.8780
Epoch 3/5
32/32 [==============================] - 0s 3ms/step - loss: 0.3037 - sparse_categorical_accuracy: 0.9270
Epoch 4/5
32/32 [==============================] - 0s 3ms/step - loss: 0.2133 - sparse_categorical_accuracy: 0.9410
Epoch 5/5
32/32 [==============================] - 0s 3ms/step - loss: 0.1592 - sparse_categorical_accuracy: 0.9700

After exporting the model, we get 2 files.

!ls tfmnist_json 
model.json group1-shard1of1.bin

The model.json will be load by tfjs library and then it will automatically download group1-shard1of1.bin.

We need to host the exported files some where so browser can download them. You can checkout my model at /tfmnist_json/model.json.

Here I use ES module feature to keep it simple. No build tools, just magic.

import * as tf from "https://cdnjs.cloudflare.com/ajax/libs/tensorflow/3.18.0/tf.fesm.js";
window.tf = tf;
const MODEL_DB = "tfmnist";
const MODEL_DB_URI = `indexeddb://${MODEL_DB}`;

let model;
try {
  model = await tf.loadLayersModel(MODEL_DB_URI);
  console.debug("Loaded model locally");
} catch (err) {
  console.error("Didn't find local model.", err);
  console.debug("Loading from remote");

  model = await tf.loadLayersModel("/tfmnist_json/model.json");
  console.debug("Loaded remote model and saving to local db.");
  model.save(MODEL_DB_URI);
}

After our model loaded successfully, we can start making predictions. Remember to preprocess the image so it can match input shape of our model.

  const image = tf.browser
    .fromPixels(context.getImageData(0, 0, canvas.height, canvas.width), 1)
    .resizeBilinear([28, 28])
    .div(255);
  const input = image.reshape([-1, 28 * 28]);
  const predictions = Object.entries(await model.predict([input]).data()).sort(
    (a, b) => b[1] - a[1]
  );

You can find the full code at main.mjs

You can also achive the same result without using Tensorflow. In fact, I trained an other model using Octave and export its weights to CSV. Loaded it in to a React app and calculate all matrix multiplications with the helps of mathjs. And it works like a charm.