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 Link to heading

Training on python Link to heading

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
from keras.models import Sequential
from keras.layers import Dense, Activation
from keras import optimizers
from datetime import datetime

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

sgd = optimizers.SGD(lr=1)
model.compile(optimizer=sgd,
              loss='binary_crossentropy',
              metrics=['accuracy'])

model.fit(X.T, Y.T, epochs=1000, batch_size=X.shape[1], verbose=0)
end_time = datetime.now()
print (end_time - start_time)

0:00:02.631679

That took 2.63s

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

1
2
3
import tensorflowjs as tfjs

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

Load the model in TensorflowJS via Javascript

1
2
3
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 Link to heading

Code used to train the model:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
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']})

train_data().then(function(){
	  console.log('Training is Complete');
	  // Do predictions and so on...
}

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

Why do we use an async function? That’s because model.fit() 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.

Conclusion Link to heading

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.