Building a Simple Machine Learning Model with Node.js, TensorFlow.js, and TypeScript

Nodejs

Node.js enables fast, scalable server-side applications with non-blocking, asynchronous execution. It runs JavaScript on the backend, making full-stack development seamless. TypeScript enhances JavaScript with static typing, improving code quality and maintainability. Together, they create a powerful, efficient, and developer-friendly environment for modern web applications.

Tensorflow.js

TensorFlow.js is a JavaScript library for building and running machine learning models directly in the browser or on Node.js. It allows developers to train and deploy AI models without needing a backend, leveraging GPUs for accelerated computations. With support for pre-trained models and custom training, it enables real-time AI applications like image recognition and natural language processing. Running in the browser enhances privacy and reduces latency by keeping data on the client side. TensorFlow.js makes machine learning accessible to web developers, enabling seamless AI integration into web applications.

Why TypeScript with TensorFlow.js?

Using TypeScript with TensorFlow.js allows machine learning models to run directly in the browser or on Node.js, enabling real-time AI without a backend. TypeScript’s static typing helps catch errors early, improving code quality and maintainability in complex ML applications. Unlike Python, which requires a server or dedicated environment, TensorFlow.js leverages JavaScript’s ecosystem for seamless frontend integration. Running models in the browser enhances privacy and reduces latency by processing data locally. This makes TypeScript and TensorFlow.js ideal for interactive, web-based AI applications without server dependencies.

Setting Up the Project
Step 1: Install Node.js and TypeScript

Make sure you have Node.js installed. You can download it from nodejs.org.

Next, install TypeScript globally:

npm install -g typescript

Step 2: Create a New Node.js Project

Create a new directory for your project and initialize it:

mkdir node-tensor
cd node-tensor
npm init -y

Step 3: Install Dependencies

Install TensorFlow.js and TypeScript dependencies:

npm install @tensorflow/tfjs-node
npm install --save-dev typescript @types/node

Step 4: Configure TypeScript

Create a tsconfig.json file to configure TypeScript:

{
  "compilerOptions": {
    "target": "ES6",
    "module": "commonjs",
    "outDir": "./dist",
    "rootDir": "./src",
    "strict": true,
    "esModuleInterop": true
  }
}

This configuration tells TypeScript to:

• Compile to ES6.

• Use CommonJS modules (for Node.js).

• Output the compiled files to the dist folder.

• Look for source files in the src folder.

Building the Model: Even/Odd Number Predictor

We’ll build a simple model to predict whether a number is even or odd. While this is a trivial problem for humans, it’s a great way to understand the basics of TensorFlow.js.

Step 1: Create the Source Files

Create a src folder and add the following files:

1. index.ts: The main entry point for our application.

2. model.ts: Contains the logic for building and training the model.

Step 2: Define the Dataset

In index.ts, let’s create a dataset of numbers and their corresponding labels (0 for odd, 1 for even):

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

// Features: Numbers from 1 to 20
const numbers = tf.tensor1d([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]);

// Labels: 0 for odd, 1 for even
const labels = tf.tensor1d([0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]);

Step 3: Build the Model

In model.ts, define a simple neural network model:

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

export function createModel(): tf.Sequential {
  const model = tf.sequential();

  // Add a single dense layer with 1 unit (neuron)
  model.add(tf.layers.dense({ units: 1, inputShape: [1], activation: 'sigmoid' }));

  // Compile the model
  model.compile({
    optimizer: tf.train.adam(0.1), // Adam optimizer
    loss: 'binaryCrossentropy',    // Binary cross-entropy for binary classification
    metrics: ['accuracy'],
  });

  return model;
}

Here’s what this code does:

• Creates a sequential model.

• Adds a single dense layer with 1 unit and a sigmoid activation function (for binary classification).

• Compiles the model using the Adam optimizer and binary cross-entropy loss.

Step 4: Train the Model

Back in index.ts, let’s train the model:

import * as tf from '@tensorflow/tfjs-node';
import { createModel } from './model';

async function trainModel() {
  const numbers = tf.tensor1d([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]);
  const labels = tf.tensor1d([0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]);

  const model = createModel();

  // Train the model
  await model.fit(numbers, labels, {
    epochs: 50, // Number of training iterations
    callbacks: {
      onEpochEnd: (epoch, logs) => {
        console.log(`Epoch ${epoch}: loss = ${logs?.loss}, accuracy = ${logs?.acc}`);
      },
    },
  });

  return model;
}

trainModel().then((model) => {
  console.log('Training complete!');

  // Test the model
  const testNumber = tf.tensor1d([21, 22, 23, 24]);
  const predictions = model.predict(testNumber) as tf.Tensor;

  predictions.print();
});

Step 5: Run the Code

Compile and run the TypeScript code:

tsc
node dist/index.js

Epoch 1: loss = 0.693, accuracy = 0.5
Epoch 2: loss = 0.692, accuracy = 0.5
...
Epoch 50: loss = 0.123, accuracy = 1.0
Training complete!
Tensor
    [0.123456, 0.987654, 0.123456, 0.987654]

The model predicts whether the numbers [21, 22, 23, 24] are even or odd. Values close to 0 indicate odd, and values close to 1 indicate even.

Explanation of the Model

1. Input: A single number (e.g., 21).

2. Output: A probability between 0 and 1 (e.g., 0.123 means the model thinks it’s odd).

3. Activation Function: The sigmoid function ensures the output is between 0 and 1.

4. Loss Function: Binary cross-entropy measures how well the model predicts binary outcomes.

5. Optimizer: Adam adjusts the model’s weights to minimize the loss.

Saving and Loading the Model

Saving the Model

To save the trained model, add this to index.ts:

await model.save('file://./model');

Loading the Model

To load the model later:

const loadedModel = await tf.loadLayersModel('file://./model/model.json');

Conclusion

In this tutorial, we built a simple machine learning model using TensorFlow.js and TypeScript in Node.js. We trained the model to predict whether a number is even or odd, and we used TypeScript to make the code more robust and maintainable.

This example is beginner-friendly and demonstrates the basics of:

• Setting up TensorFlow.js with TypeScript.

• Creating and training a simple neural network.

• Making predictions with the trained model.

You can extend this example by:

• Adding more layers to the model.

• Using a larger dataset.

• Experimenting with different activation functions and optimizers.

Happy coding!