Tensorflow.js: Machine Learning through JavaScript

Sign up for FREE 1 month of Kindle and read all our books for free.

Get FREE domain for 1st year and build your brand new site

Reading time: 30 minutes | Coding time: 15 minutes

Tensorflow.js is an open-source library (for TensorFlow framework) with which we can implement Machine Learning in the browser with the help of JavaScript. We can implement many models of machine learning through JavaScript like webcam controller, teachable machine, Node.Js pitch prediction, and many more.

Tensorflow.js supports WebGL and it uses GPU acceleration (means it will automatically accelerate the code in the backend when the GPU is available). It is useful for privacy preserved applications as well as low-latency inference since all the data stays on the client.

TensorFlow.js is powered by WebGL and provides a high-level layers API for defining models, and a low-level API for linear algebra and automatic differentiation. TensorFlow.js also supports importing TensorFlow SavedModels.

Workflow for Tensorflow.js

  1. An existing, pre-trained model can be imported for inference -> If we have any existing, previously worked model for tensorflow then it can be converted to tensorflow.js format and can be loaded to the browser for inference.

  2. An existing, previous model can be rebuild or can be worked upon -> Since, We can import the existing model therefore we can also work on it to rebuild it.

  3. Use models directly into the browser -> As discussed earlier, we can implement models through javascript and run it directly in the browser.

Coding Section

Let us see, how to import Python defined models for machine learning in the browser as well as how to define it using javascript. As an example, here is the code that defines a neural network to classify flowers:

import * as tf from ‘@tensorflow/tfjs’;
const modelImported = tf.sequential();
modelImported.add(tf.layers.dense({inputShape: [4], units: 100}));
modelImported.add(tf.layers.dense({units: 4}));
modelImported.compile({loss: ‘categoricalCrossentropy’, optimizer: ‘sgd’});

The model for the network is:

const input = tf.tensor2d([[4.8, 3.0, 1.4, 0.1]], [1, 4]);
const result = model.predict(inputData);
const winner = irisClasses[result.argMax().dataSync()[0]];

In the above code, "input" is for getting the measurements for the new flower for generating a prediction. The First argument is for the data and second is for shape.
Now, "result" and "winner" is for getting the hoghest confidence prediction from the model. And finally, displaying it in the console.

Running the code

We can run the program in two different ways. Firstly, we can add the CDN of tensorflow.js in "head" section of HTML file with the help of "script" tags. For example:

        <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@0.12.0"> </script>

Or, we can install it via yarn or NPM ->

npm install @tensorflow/tfjs
yarn add @tensorflow/tfjs

After installing TensorFlow in your project, you can import it to main file by the following code:

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

Core APIs of TensorFlow.js


A tensor is a n-dimensional array with n>2. Tensorflow.js has utility functions for scaler, 1D, 2D, 3D, 4D tensors and it also has variety of functions to initialize tensors in ways that are useful for machine learning.


Tensors show the property of immutability. Immutability means that the value of tensors once given cannot be changed. Therefore, tf.variable() is introduced in tensorflow.js. The use for tf.Variable() is when we need to change the data frequently such as for adjusting model weights. "Variable()" requires an initial value which can be a Tensor of any type and shape. This value defines the type and shape of the variable.

However, if you want to change the value of the variable, it can be changed by .assign() method. Example:

const x = tf.variable(tf.tensor([1, 2, 3]));
x.assign(tf.tensor([4, 5, 6]));

Mathematical Operations in Tensorflow.js

To perform Mathematical computations in tensorflow.js, we use operations. Since, tensors are immutable, therefore, the operations always return new tensors and the input Tensors are never modified.

Lets see some of the operations:


It adds two tf.tensors() element wise. For example ->

const x = tf.tensor1d([5, 6, 7, 8]);
const y = tf.tensor1d([10, 20, 30, 40]);
tf.add(x, y).print();  


It computes the dot product of two matrices. For example ->

const a = tf.tensor2d([1, 2], [1, 2]);
const b = tf.tensor2d([1, 2, 3, 4], [2, 2]);
tf.matMul(a, b).print();


It subtracts two tensors element wise. For example ->

const a = tf.tensor1d([10, 20, 30, 40]);
const b = tf.tensor1d([1, 2, 3, 4]);
tf.sub(a, b).print();


It multiplies two tensors element wise. For example ->

const a = tf.tensor1d([1, 2, 3, 4]);
const b = tf.tensor1d([2, 3, 4, 5]);
tf.mul(a, b).print();


It divides two Tensors element-wise.

const a = tf.tensor1d([1, 4, 9, 16]);
const b = tf.tensor1d([1, 2, 3, 4]);
tf.div(a, b).print();  

And many more functions like these.

Memory Management in TensorFlow.js

TensorFlow.js provides two major ways to manage memory:

  1. tf.dispose()
  2. tf.tidy()


To destroy the memory of a tf.Tensor, we can use the tf.dispose() methos. It disposes any Tensors found within the mentioned object. For example ->

const dis= tf.scalar(2);


This method cleans up all tf.Tensors that are not returned by a function after executing it, similar to the way local variables are cleaned up when a function is executed. For example ->

const x = tf.tensor([[1, 2], [3, 4]]);
const y = tf.tidy(() => {
  const result = x.square().log().neg();
  return result;


Layers are the primary building block for constructing a machine learning model. Each layer performs some computation to calculate its output.
Layers automatically create and initialize the variables that are needed to function, increasing the level of abstraction. Now, let us take a look on the code:

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

// Build and compile model.
const mlModel = tf.sequential();
mlModel.add(tf.layers.dense({units: 1, inputShape: [1]}));
mlModel.compile({optimizer: 'sgd', loss: 'meanSquaredError'});

// Generate some synthetic data for training.
const xs = tf.tensor2d([[1], [2], [3], [4]], [4, 1]);
const ys = tf.tensor2d([[1], [3], [5], [7]], [4, 1]);

// Train model with fit().
await model.fit(xs, ys, {epochs: 1000}); //epochs refers to the iterations in a dataset.

// Run inference with predict().
mlModel.predict(tf.tensor2d([[5]], [1, 1])).print();

With this article at OpenGenus, you have a good understanding of the basics of TensorFlow.js and you will be able to proceed further and use Machine Learning in JavaScript. Enjoy.