By guest author Aayush Arora (GitHub :

This codelab will walk you through creating your own neural network, using TensorFlow.js, and train it using the Iris Flower dataset, and then categorize the dataset into three classes.

What you will build

In this codelab, you're going to build a simple neural network that can classify irises into three classes using TensorFlow.js:

  1. Setosa
  2. Virginica
  3. Versicolor

Dataset Used: Iris

What you'll learn

What you'll need

This codelab is focused on TensorFlow.js. Less relevant concepts and code blocks are glossed over and are provided for you to simply copy and paste.

Download the Code

Click the following link to download all the code for this codelab:

Clone all the code for this codelab:

git clone

Or Download Source Code

This will provide you the repository TensorflowJS-NeuralNet, which contains the folder challenge with file index.js, and the dataset iris.json. We'll be doing all our coding work in the file called index.js.

Running the project

To run the project, you need to run the command npm install from the challenge directory.

Once the packages are successfully installed, you will be able to see this:

How to convert the data into the TensorFlow.js format?

The tf.tensor2D function helps in creating data that TensorFlow.js understands well. As our dataset is a flat array, we will need to pass the shape as the second parameter to this function. Our first step will be to create our training and testing dataset.

Creating training and testing data

There are multiple ways to get started with any project. In this case, to keep our project as simple as possible and concentrate on a Neural Network, we've provided you with basic code along with the dataset.

We will divide the dataset into two parts: trainingData and testingData. To divide the iris.json into these two parts, such that they don't have any intersection, run this command in the challenge folder:

node generateTest.js

The command will divide the 144 total data items into 130 training data items and 14 testing data items.

The trainingData will have the shape of [130,4]as the length of training dataset array is 130 and there are four features associated with each object. Insert this code into your challenge/index.js file:

const trainingData = tf.tensor2d(> [
        item.sepal_length, item.sepal_width, item.petal_length, item.petal_width

You need to create the testingData similarly, using testing.json.

Declare a variable testingData,think about its shape, and write the code for the testing dataset in challenge/index.js.

 // Your code here

The output will be based on the neuron activation. We will write the output function such that we get an array of length three every time with one of the values closer to one and the rest two closer to zero.

const outputData = tf.tensor2d( => [
    item.species === 'setosa' ? 1 : 0,
    item.species === 'virginica' ? 1 : 0,
    item.species === 'versicolor' ? 1 : 0

]), [130,3])

This procedure consists of two steps:

Creating Model

Creating a model in TensorFlow.js is super easy with a single line of code:

const model = tf.sequential();

Adding Input Layer

Here, we need to add the Input Layer to our model. Since, we have four features, we will use 4 input nodes.

    inputShape: [4],
    activation: "sigmoid",
    units: 10

Adding Output Layer

// based on the input layer definition above, add an output
// layer to challenge/index.js

Compiling your Model

We need to use model.compile function with an optimizer to decrease the loss.

    loss: "categoricalCrossentropy",
    optimizer: tf.train.adam()

Fitting and Predicting your Model

Now, you can feed your model with the training and OutputData and predict new Data.

You can also validate your results using your testingData.

Finally, you can print the results in the form of an Array on the screen using


async function train_data(){
   for(let i=0;i<15;i++){
      const res = await,
                  outputData,{epochs: 40});  

Now, you can write the main function which will call train_data and will predict the output after train_data will complete its execution.

async function main() {
  let train = await train_data();

After training inside the main function, to predict the result, we will use the model.predict function with testingData as a parameter.

Finally, we will print the result using the .print() function.

async function main() {
  await train_data();
  // Write your predict function here and print the results

After you have completed this code, you can directly run the program like this:

node index.js

The model will classify the testingData into the matrices of the form [1,0,0], [0,1,0] or [0,0,1] which directly relates as follows:

As per our OutputData function:

[1,0,0] corresponds to setosa

[0,1,0] corresponds to verginica

[0,0,1] corresponds to versicolor

When we run this file with the given testing dataset, we should see the results as follows with this loss history:

This clearly indicates our testing data contains versicolor, verginica and setosa.