Welcome to the Machine: a gentle introduction to Machine Learning in JavaScript 🤖

The hot topic in tech right now is most certainly AI and everything that’s involved in it. Machine learning is perhaps the most important aspect for AI but, as JavaScript developers, is it something we can do in our beloved programming language?

The title of this article sort of kills the suspense and so yes! It’s possible! Even better news is that this article is going to provide a gentle introduction to machine learning in JavaScript.

We’ll walk through the process of creating a model that can make predictions based on data we provide, and I’ll explain each step so that you don’t get lost along the way.

It’s a gentle introduction, so the aim is to help you understand how you can implement machine learning in JavaScript without going too deep into the machine learning abyss.

Photo by Shane Aldendorff: https://www.pexels.com/photo/flat-lay-photography-of-black-and-gray-components-on-white-surface-924675/

Ok, first thing’s first. There are quite a few good JS machine learning libraries out there, and one of the most popular ones is TensorFlow.js.
This is the library we will be using to create our prediction model.

This walkthrough is meant for JS developers of all levels and so feel free to skip the parts you know already!

Setting up your project

We’re going to start from absolute scratch, so what you’ll need to do is create a folder and give it a name. Mine is called machine-learning-beginners. Open this folder in your code editor of choice — I will be using VS Code- open the terminal and make sure you are in the correct folder:

Initial steps

We then need to run npm init to create a Node.js project. Once you run this command in your terminal you will be prompted to add some information about the project. I’m going to be super lazy and I’m just going to leave it empty for now (you can either add the info now or later but it’s not too important for the demo):

npm init setup

We then need to install TensorFlow.js:

npm install @tensorflow/tfjs

Finally, create an index.js file and import TensorFlow.js:

// index.js

// Use CommonJS require syntax to import TensorFlow.js
const tf = require("@tensorflow/tfjs");

Setup done! Now, we can start writing some machine learning code!

Creating a prediction model in JavaScript

First it’s probably good to outline what our model should do:

If you’ve read my previous articles you might know that I’m a huge football (the round one) fan, and so we are going to create a model that predicts the result of a football match based on the previous results between two teams.

Training Data
Now we can get started on the code. 😏

The most important part of machine learning is data, and so we need to provide our model with data to train it. The data we are going to provide is the last 20 results between Liverpool FC and Everton FC. The focus will be on Liverpool and so they will always be the first input:

// index.js

const tf = require("@tensorflow/tfjs");
// Step 1: Add the training data
// Liverpool FC = 0
// Everton = 1
const trainingData = [
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [0, 1, 0] }, // Liverpool lose
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool win
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [0, 0, 1] }, // Draw
];

Ok… so what does all that mean?

Well, first of all we have created an array of objects, and each object shows one game result.

Inputs
We have two inputs: 0 and 1.
0 represents Liverpool and 1 represents Everton.

Remember, we always have Liverpool as the first input because our focus is on Liverpool’s results.

We have to represent our categorical data (the football team names) using numerical values because machine learning algorithms primarily work with numerical data. But don’t let that confuse you.

Outputs
We have three outputs which represent: win, lose, draw. If Liverpool win, the output will be 1, 0, 0. If they lose it will be 0 ,1, 0. If they draw it will be 0, 0, 1.

Nice! It’s not a lot of training data but it’s good enough for our first prediction model. So, let’s dive in:

// index.js

// ... add this after the trainingData array
// Step 2: Construct the model
const model = tf.sequential();
model.add(tf.layers.dense({ units: 10, inputShape: [2], activation: "relu" }));
model.add(tf.layers.dense({ units: 3, activation: "softmax" }));

Here is the same code with added context:

// index.js

// ... add this after the trainingData array
// Step 2: Construct the model
const model = tf.sequential(); 
// creates a model where each layer is connected
// to the next in a sequential manner 
model.add(tf.layers.dense({ units: 10, inputShape: [2], activation: "relu" }));
// This is a hidden layer. 
// Allows the model to learn more complex representations 
// and extract meaningful features from the input data.
// inputShape: [2] indicates our 2 inputs: Liverpool and Everton.
// relu is used for introducing non-linearity in neural networks.

model.add(tf.layers.dense({ units: 3, activation: "softmax" }));
// This is the output layer. 
// Generates probabilities for each outcome, 
// indicating the likelihood of each outcome occurring.
// units: 3 represents the probability our 3 outputs (win, lose, draw).
// Softmax activation function is used for multi-class classification problems
// as it converts the output values into probabilities that sum up to 1.

In the code above, we create a (neural network) model and then create 2 (dense) layers to help the model learn about the patterns in the data.

If this doesn’t make sense to you, don’t worry — we’re basically just helping the model to learn about the shape of our training data!

Adding multiple layers to a neural network allows it to learn more complex patterns and relationships in the data. Each layer performs specific computations and transformations on the input data, gradually learning patterns and making predictions.

Photo by Giant Asparagus: https://www.pexels.com/photo/top-view-of-a-neuron-sculpture-in-the-lawn-9599917/

Now we need to compile the model, which optimises it for training:

// index.js 

// ... add this after previous code 
// Step 3: Compile the model
model.compile({ loss: "categoricalCrossentropy", optimizer: "adam" });

With added context:

// index.js 

// Step 3: Compile the model
model.compile({ loss: "categoricalCrossentropy", optimizer: "adam" });
// Categorical cross-entropy measures the dissimilarity between the predicted
// probabilities and the true labels, which makes the model output higher
// probabilities for the correct class.
// adam is an optimization algorithm that adapts the learning rate during 
// training and combines the benefits of both AdaGrad and RMSProp optimizers. 
// It is widely used for training neural networks and helps in efficiently 
// updating the model's parameters to minimize the loss.

We have prepared our model for the training data, and now we also need to prepare trainingData for the model:

// index.js 

// ...add this after previous code
// Step 4: Prepare the training data
const xTrain = tf.tensor2d(trainingData.map((item) => item.input));
const yTrain = tf.tensor2d(trainingData.map((item) => item.output));

In the above code, we create two variables: xTrain and yTrain, which we will use to train our model. In machine learning, it is common to represent data in a structured format that can be easily processed by the learning algorithm. For neural networks, the data should be in the form of tensors.

This is a ‘gentle’ introduction so we won’t go into the mathematics behind this, but you can use console.log to check the respective values if this is of interest to you.

Now that we have prepared our training data and model, we can finally use our data to train the model:

// index.js 

// ...add this after previous code
// Step 5: Train the model
async function trainModel() {
  await model.fit(xTrain, yTrain, { epochs: 100 });
  console.log("Training complete!");
}
// epochs: 100 - the model will iterate over the entire training dataset 100
// times and will update its internal parameters (weights and biases) 
// in each iteration. This will minimise the loss and improve its performance.

Note: Increasing the number of epochs can potentially improve the model’s performance, as it allows the model to learn from the data more thoroughly.

However, in setting a larger number of epochs you may risk overfitting the model to the training data.

The optimal number of epochs depends on the specific problem and dataset and is often determined through experimentation and validation.

Woo! We now have a trained model that should be able to predict the outcome of Liverpool’s next game against Everton.

All we need now is to provide test data and let our model start to make predictions:

// index.js 

// ...add this after previous code
// Step 6: Make predictions
function makePredictions() {
  const testData = [
    { input: [0, 1] } // Liverpool vs. Everton
  ];
  const predictions = model.predict(
    tf.tensor2d(testData.map((item) => item.input))
  );
  const predictedResults = Array.from(predictions.dataSync());
//  uses the trained model to make predictions on the test data

  testData.forEach((data, index) => {
    const predictedOutcome = predictedResults.slice(index * 3, (index + 1) * 3);
    console.log(
      `Match ${index + 1}: Liverpool ${data.input[0]} vs. Everton ${
        data.input[1]
      }`
    );
// extracts the predicted outcome for our test case
// Since each prediction consists of three values (win, lose, draw),
// slice() is used to extract the relevant values for the current test case.
    console.log("Predicted outcome:", predictedOutcome);
  });
}

The function above takes test data (e.g., { input: [0, 1] // Liverpool vs. Everton), makes predictions using the trained model, and then logs the predicted outcome of a single match between Liverpool and Everton.

The last thing we need to do is put all of this together:

// index.js 

// ...add this after previous code
// Step 7: Train the model and make predictions
trainModel().then(() => {
  makePredictions();
});

In the code above, we first train the model and then use it to make predictions.

Your final index.js file should look like this (comments removed):

const tf = require("@tensorflow/tfjs");

// Liverpool FC = 0
// Everton = 1
const trainingData = [
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [0, 1, 0] }, // Everton won
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [1, 0, 0] }, // Liverpool won
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [0, 0, 1] }, // Draw
  { input: [0, 1], output: [0, 0, 1] }, // Draw
];
const model = tf.sequential();
model.add(tf.layers.dense({ units: 10, inputShape: [2], activation: "relu" }));
model.add(tf.layers.dense({ units: 3, activation: "softmax" }));
model.compile({ loss: "categoricalCrossentropy", optimizer: "adam" });
const xTrain = tf.tensor2d(trainingData.map((item) => item.input));
const yTrain = tf.tensor2d(trainingData.map((item) => item.output));
async function trainModel() {
  await model.fit(xTrain, yTrain, { epochs: 100 });
  console.log("Training complete!");
}
function makePredictions() {
  const testData = [
    { input: [0, 1] }, // Liverpool vs. Everton
  ];
  const predictions = model.predict(
    tf.tensor2d(testData.map((item) => item.input))
  );
  const predictedResults = Array.from(predictions.dataSync());
  testData.forEach((data, index) => {
    const predictedOutcome = predictedResults.slice(index * 3, (index + 1) * 3);
    console.log(
      `Match ${index + 1}: Liverpool ${data.input[0]} vs. Everton ${
        data.input[1]
      }`
    );
    console.log("Predicted outcome:", predictedOutcome);
  });
}
trainModel().then(() => {
  makePredictions();
});

So… does it work?

Run node and the name of your file in the terminal.
For example: node index.js

output of node index.js

Note that we trained our model with a very small amount of data and so the predictions it provides may vary a bit. But, considering Liverpool have won the majority of their last matches against Everton, it seems our model is working pretty well!

To confirm that it’s working, I will change trainingData to say that Liverpool won all of their last matches against Everton (hehe).

Now let’s run node index.js again:

output of node index.js

A 91% chance of winning seems pretty sweet to me. So that’s it! That’s our ‘gentle’ introduction to machine learning in JavaScript. YNWA 🤖

I hope you enjoyed it and please leave feedback if there are any parts which were unclear, as I would love for this article to be understandable to developers of all levels.