To Tensorflow JS or not?

I was writing a blog post about the different activation functions and how they can impact your training. I wanted to make an interactive post so was looking at running a simple model on the browser. That’s when I found Tensorflow JS.

Kudos to the people behind this. This was a great way to bring the awesomeness of Tensorflow to the browser. When using the library you can either:

  1. Load a pre-trained model
  2. Train on the browser

Let’s try and do both. We’ll work with the make_moons dataset from sklearn.

Parameters we’ll be using:

  1. Activation Function: tanh
  2. Learning Rate: 1

The data we’ll be training on:

Load a pre-trained model

Training on python

from keras.models import Sequential
from keras.layers import Dense, Activation
from keras import optimizers
from datetime import datetime

start_time =
model = Sequential()
model.add(Dense(4, activation='tanh', input_dim=X.shape[0]))
model.add(Dense(1, activation='sigmoid'))

sgd = optimizers.SGD(lr=1)
              metrics=['accuracy']), Y.T, epochs=1000, batch_size=X.shape[1], verbose=0)
end_time =
print (end_time - start_time)


That took 2.63s

Save the model to a form which can be used by TensorflowJS

import tensorflowjs as tfjs

tfjs.converters.save_keras_model(model, "/tmp/")

Load the model in TensorflowJS via Javascript

import * as tf from '@tensorflow/tfjs';

const model = await tf.loadModel("/tensorflowjs/model.json");

Finally Lets see how our model predicts

Train on the browser

Code used to train the model:

const model = tf.sequential()
model.add(tf.layers.dense({units:4, inputShape:2, activation:'tanh'}))
model.add(tf.layers.dense({units:1, activation:"sigmoid"}))
var sgd = tf.train.sgd(1)
model.compile({optimizer:sgd, loss:'binaryCrossentropy', metrics:['accuracy']})

	  console.log('Training is Complete');
	  // Do predictions and so on...

async function train_data(){
	const res = await, Y, {epochs:1000, batchSize:200, verbose: 1})

Why do we use an async function? That’s because returns a Javascript promise. Hence we start an async function, and wait on it to finish.

Time taken to finish training: Loading

If you wait long enough, you’ll see the model finish training and prediction similar to our earlier method. The only difference is that this time it took a lot more time. On my PC it takes around 22s which is almost 1000% of the previous method.


Its amazing that we can even do something on the lines of training models on browsers. But at the same time, we can see how years of optimization on the Python front of Tensorflow clearly gives it an edge over its JS counterpart.