Starting with TensorFlow.js in Node.js

Hello Readers. In this blog, I will show you how to use the TensorFlow.js library with Node.js as the backend. You must all wonder that to connect your deep learning models with your development would be such a massive task. If you are a developer and use Node.js for the backend, this tutorial will help you set up the TensorFlow.js library in Node.js. Now you don’t have to code in Python to deploy your Machine and Deep Learning models to the browser.

For your introduction – TensorFlow.js is an open-source JavaScript library for training and deploying ML/DL models.

To get TensorFlow.js in your web-projects:

Using a script tag:

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/[email protected]/dist/tf.min.js"></script>

Using the command-line interface of NPM:

npm install @tensorflow/tfjs

After you are done with the installation part, you can import different types of TensorFlow packages in your backend through Node.

Types of TensorFlow Packages

TensorFlow CPU

import * as tf from '@tensorflow/tfjs-node'

If you import the TensorFlow CPU package then all the algebraic computations of the library would run on CPU.

TensorFlow GPU

import * as tf from '@tensorflow/tfjs-node-gpu'

On the other hand, importing the GPU package will run all the tensor operations on the GPU with CUDA. To work with GPU you need to have an NVIDIA Graphics card with CUDA installed on your computer.

Vanilla CPU

import * as tf from '@tensorflow/tfjs'

And lastly, you can simply use the above package in which operations are run in vanilla JavaScript on the CPU. This is a small package and is most compatible as it can be used in more devices that support Node.js than just Linux, Windows. But this package takes more time for computations and is slower in speed than the above two.

In the below code we have implemented TensorFlow.js to build a simple Linear Regression Model.

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

//Simple Linear Regression model
const model = tf.sequential();
//only 1 input layer of unit 1
model.add(tf.layers.dense({units: 1, inputShape: [1]}));

//Specifing the loss and the optimizer 
model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});

//Giving some self made data points
const xaxis = tf.tensor2d([1, 2, 3, 4], [4, 1]);
const yaxis = tf.tensor2d([2, 4, 6, 8], [4, 1]);

// Training and prediction
model.fit(xaxis, yaxis).then(() => {
  model.predict(tf.tensor2d([5], [1, 1])).print();
});

Firstly, we import TensorFlow.js and then describe a model as a constant.
We add only 1 hidden layer for simplicity and just compile the model with the specified loss and opmitizer.
Then, the input data points are created for the model and passed for the training of the model.
Lastly, the value is predicted for input = 5. We get 9.89 ~ 10 as our output.