Simple neural network in JavaScript using TensorFlow.js

In this tutorial, I will let you know how to create your first simple artificial neural network using the JavaScript TensorFlow.js library.

The main purpose of this article is to give you the simplest example of using the TensorFlow.js machine learning library. The example of the neural network we are going to give you here is actually a linear regression and can be run on the browser.

Here we will predict the price for a given number of rooms from our already given data.

Complete JavaScript code to build a simple neural network in TensorFlow.js

Before I go through the explanation, see the complete JavaScript code below for TensorFlow.js price prediction for a given number of room:

// Create our async function
async function runNN()
{

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

// training datasets
// In our training datasets, we take room numbers and corresponding price to rent
const xs = tf.tensor2d([
     [2],[4],[3],[8],[5],[6],
  ]);
const ys = tf.tensor2d([
     [1900],[3740],[2450],[7324],[4200],[5600],
  ]);

// Train model with fit().method
await model.fit(xs, ys, {epochs: 1000});

// Run inference with predict() method.
model.predict(tf.tensor2d([
     [6],[3],[8]
  ])).print();

}

// Call our function
runNN();

That’s it… This is our complete code for the simplest neural network in TensorFlow.js which you can run on your browser. You can only see the input and output node and there are no hidden layers which makes it so simple.

Also, read:

In our code, carefully read the comments to understand what we did. Be sure, that you include the TensorFlow.js JavaScript library. You can include it from CDN.

Breaking our code into pieces for better understanding

Now let me break our code into pieces and describe briefly.

First of all, we have created our async function:

async function runNN()
{
   // All our code will be here and later we will call this function
}

We have defined the function async because we will use await while training our model. The await operator only can work inside the async function.

Now, let’s start coding inside our async function…

First of all, create our model object:

const model = tf.sequential();

As you can see in the above code, creating model in TensorFlow.js is very easy.

Now add our input layer to the model we created:

model.add(tf.layers.dense({
  units: 1,
  inputShape: [1]
}));

As we have only one feature, so here we have added input node 1 which passes through inputShape item.

Next, compile our model with the compile() method:

model.compile({
  optimizer: 'sgd',
  loss: 'meanSquaredError'
});

Now let’s take our training datasets:

// training datasets
// In our training datasets, we take room numbers and corresponding price to rent
const xs = tf.tensor2d([
     [2],[4],[3],[8],[5],[6],
  ]);
const ys = tf.tensor2d([
     [1900],[3740],[2450],[7324],[4200],[5600],
  ]);

In our training dataset, we have the number of rooms as features and prices as the label.

Now it’s time to train our model with our training data using the fit() function:

// Train model with fit().method
await model.fit(xs, ys, {epochs: 1000});

Here we have set epochs to 1000. The higher value of epochs will give more accurate results, but it will take more processing time. So it should be taken wisely. With epochs, we set how many times our training datasets will be taken to update our model.

Now, we can get a prediction for new data as well as data already exist in training to predict the result:

// Run inference with predict() method.
model.predict(tf.tensor2d([
     [6],[3],[8]
  ])).print();

}

Finally, we can call our function and get the result in the console:

runNN();

We can see tensor in our console that contains predicted price for our given set of the number of rooms in the predict function. We can see it predicting the almost meaningful price that matches with our training datasets.

We can give any set of data to predict and it will give us the result when we run it on our browser.

 

Leave a Reply

Your email address will not be published. Required fields are marked *