HomeAudienceBeginnersBeginning Machine Learning with TensorFlow.js

Beginning Machine Learning with TensorFlow.js

- Advertisement -

With artificial intelligence (AI) taking centre stage in the deployment of data, machine learning (ML), which is an important application of AI, gains importance. ML is much helped by the use of TensorFlow.js, an open source javaScript library this is used to define, train and run ML models.

Machine learning is an application of artificial intelligence (AI) that provides systems with the ability to automatically learn and improve from experience, without being explicitly programmed. Machine learning (ML) focuses on the development of computer programs that can access data and use it to learn for themselves.

Basically, ML systems process the data and the output to come up with rules that match the input data to the output.

- Advertisement -


TensorFlow is an open source software library for numerical computation and it uses data flow graphs. The graph nodes represent mathematical operations, while the graph edges represent the multi-dimensional data arrays (tensors) that flow between them.

TensorFlow was originally developed by researchers and engineers working at Google. The system is general enough to be applicable to problems related to a wide variety of domains.


Tensorflow.js is an open source library that makes use of JavaScript and a high-level layer API to define, train as well as run machine learning models entirely in the browser. This open source library is powered by WebGL, which provides a high-level layer API for defining models, and offers a low-level API for linear algebra and automatic differentiation.

What can be done with TensorFlow.js

Listed below are some of the capabilities of TensorFlow.js.

  • The ability to import an existing, pre-trained model or inference: If the user has an existing TensorFlow or Keras model that has been previously trained offline, it can be converted into the TensorFlow.js format and then loaded into the browser for inference.
  • The ability to retrain an existing model: By making use of a small amount of data collected in the browser using a technique called image retraining, you can make use of transfer learning in order to augment an existing model that has previously been trained offline. This is one way to train an accurate model quickly, using only a small amount of data.
  • The ability to author models directly in the browser: You can also use TensorFlow.js with JavaScript and a high-level layer API to define, train as well as run models entirely in the browser.

Benefits of TensorFlow.js

Listed below are some of the benefits of TensorFlow.js:

  • It supports WebGL, out-of-the-box.
  • It supports GPU acceleration and therefore, behind the scenes, will accelerate your code when a GPU is available.
  • All data stays on the client, which makes TensorFlow.js useful not only for low-latency inference but for privacy-preserving applications as well.
  • Even from a mobile device, you can open your Web page, in which case your model can take advantage of sensor data.

Core concepts of TensorFlow.js

TensorFlow.js provides low-level building blocks for machine learning as well as a high-level, Keras-inspired API for constructing neural networks. Let’s take a look at some of the core components of the library.

Tensors: The central unit of data in TensorFlow.js is the tensor—a set of numerical values represented as an array of one or more dimensions. A tensor also contains the shape attribute, which defines the array’s shape and essentially gives the count of elements per dimension.

The most common way to create a tensor is by using the tf.tensor function, as shown in the code snippet below:

// 2x3 Tensor

const shape = [2, 3]; // 2 rows, 3 columns

const a = tf.tensor([2.0, 3.0, 4.0, 20.0, 30.0, 40.0], shape);

a.print(); // print Tensor values

// Output: [[2,3,4 ],

// [20, 30, 40]]

However, for constructing low-rank tensors, Google advises using the tf.scalar, tf.tensor1d, tf.tensor2d, tf.tensor3d and tf.tensor4d functions.

The following code example creates an identical tensor to the one shown in the previous code snippet using tf.tensor2d:

const c = tf.tensor2d([[2.0, 3.0, 4.0], [20.0, 30.0, 40.0]]);


// Output: [[2,3,4 ],

// [20, 30, 40]]

TensorFlow.js also provides functions for creating tensors with all values set to 0 (tf.zeros) or all values set to 1 (tf.ones).

// 3x5 Tensor with all values set to 0

const zeros = tf.zeros([3, 5]);

// Output: [[0, 0, 0, 0, 0],

// [0, 0, 0, 0, 0],

// [0, 0, 0, 0, 0]]

Variables: Variables are initialised with a tensor of values. Unlike tensors, however, variables are mutable. You can assign a new tensor to an existing variable using the assign method, as shown below:

const initialValues = tf.zeros([6]);

const biases = tf.variable(initialValues); // initialize biases

biases.print(); // output: [0, 0, 0, 0, 0,0]

const updatedValues = tf.tensor1d([0, 1, 0, 1, 0,1]);

biases.assign(updatedValues); // update values of biases

biases.print(); // output: [0, 1, 0, 1, 0,1]

Operations: While tensors allow you to store data, operations (ops) allow you to manipulate that data. TensorFlow.js provides a wide array of operations suitable for general computations and ML-specific operations on tensors. Because tensors are immutable, operations return new tensors after computations are done—for example, a unary operator such as a square.

const data = tf.tensor2d([[2.0, 3.0], [4.0, 5.0]]);

const data_squared = data.square();


// Output: [[4, 9 ],

// [16, 25]]

So is the case with binary operations such as add, sub and mul, as shown below:

const e = tf.tensor2d([[1.0, 2.0], [3.0, 4.0]]);

const f = tf.tensor2d([[5.0, 6.0], [7.0, 8.0]]);

const e_plus_f = e.add(f);


// Output: [[6 , 8 ],

// [10, 12]]

TensorFlow.js has a chainable API; you can call operations on the result of operations.

const sq_sum = e.add(f).square();


// Output: [[36 , 64 ],

// [100, 144]]

All operations are exposed as functions in the main name space; so you could also do the following:

const sq_sum = tf.square(tf.add(e, f));

Models and layers: Conceptually, a model is a function that takes some input and produces some output. In TensorFlow.js, there are two ways to create models. You can use operations directly to represent the work of the model. For example:

// Define function

function predict(input) {

// y = a * x ^ 2 + b * x + c

// More on tf.tidy in the next section

return tf.tidy(() => {

const x = tf.scalar(input);

const ax2 = a.mul(x.square());

const bx = b.mul(x);

const y = ax2.add(bx).add(c);

return y;



// Define constants: y = 2x^2 + 4x + 8

const a = tf.scalar(2);

const b = tf.scalar(4);

const c = tf.scalar(8);

// Predict output for input of 2

const result = predict(2);

result.print() // Output: 24

Or you can use APIs like tf.model and tf.sequential to construct a model. The following code constructs a tf.sequential model:

const model = tf.sequential();



units: 20,

recurrentInitializer: ‘GlorotNormal’,

inputShape: [80, 4]



const optimizer = tf.train.sgd(LEARNING_RATE);

model.compile({optimizer, loss: ‘categoricalCrossentropy’});

model.fit({x: data, y: labels});

There are many kinds of layers available in TensorFlow.js. A few examples include tf.layers.simpleRNN, tf.layers.gru and tf.layers.lstm.

Memory management

TensorFlow.js uses the GPU to accelerate math operations, so it’s necessary to manage GPU memory while working with tensors and variables.

TensorFlow.js provides two functions for this.

Dispose: You can call dispose on a tensor or variable to purge it and free up its GPU memory:

const x = tf.tensor2d([[0.0, 2.0], [4.0, 6.0]]);

const x_squared = x.square();



Tf.tidy: Using dispose can be cumbersome when doing a lot of tensor operations. TensorFlow.js provides another function, tf.tidy, which plays a similar role to regular scopes in JavaScript, but for GPU-backed tensors.

tf.tidy executes a function and purges any intermediate tensors created, freeing up their GPU memory. It does not purge the return value of the inner function.

tf.tidy takes a function to tidy up after:

const average = tf.tidy(() => {

tf.tidy will clean up all the GPU memory used by tensors inside this function, other than the tensor that is returned.

Even in a short sequence of operations like the one below, a number of intermediate tensors get created. So it is a good practice to put your math ops in a tidy!

const y = tf.tensor1d([1.0, 2.0, 3.0, 4.0]);

const z = tf.ones([4]);

return y.sub(z).square().mean();


average.print() // Output: 3.5

Using tf.tidy will help prevent memory leaks in your application. It can also be used to more carefully control when memory is reclaimed. The function passed to tf.tidy should be synchronous and also not return a promise.

tf.tidy will not clean up variables. As variables typically last through the entire life cycle of a machine learning model, TensorFlow.js doesn’t clean them up even if they are created in a tidy; however, you can call dispose on them, manually.

- Advertisement -
Abhinav Nath Gupta
Abhinav Nath Gupta
The author is a software development engineer at Cleo Software India Pvt Ltd, Bengaluru. He is interested in cryptography, data security, cryptocurrency and cloud computing. He can be reached at abhi.aec89@gmail.com.


Please enter your comment!
Please enter your name here

Thought Leaders

Open Journey

- Advertisement -