In this post, we'll be looking at how to build end-to-end machine learning workflows with TensorFlow.js and Node.js.
TensorFlow.js is a powerful toolkit for creating and training machine learning models in the browser. Node.js is a JavaScript runtime that allows you to run JavaScript code outside of the browser.
Together, these two technologies allow you to build complete machine learning workflows that can be deployed anywhere.
Before we dive into the details of how to build end-to-end machine learning workflows, let's take a moment to set up our development environment.
You'll need to have Node.js and npm installed on your machine. You can find instructions for doing so here:
https://nodejs.org/en/download/
Once you have Node.js and npm installed, you can install the TensorFlow.js CLI tool by running the following command:
npm install -g @tensorflow/tfjs-node
Now that we have our environment set up, let's create a new project. We'll start by initializing a new npm project:
npm init
This will create a file called package.json
in your project directory. This file contains information about your project, including the dependencies that your project needs.
Next, we'll install the TensorFlow.js and Node.js dependencies that we'll need for our project:
npm install @tensorflow/tfjs @tensorflow/tfjs-node
Now that we have our project set up, let's start building our machine learning model.
We'll start by creating a file called model.js
in our project directory. This file will contain the code for our machine learning model.
First, we'll need to load the TensorFlow.js and Node.js libraries that we'll be using:
const tf = require('@tensorflow/tfjs');
const fs = require('fs');
Next, we'll define some constants that we'll use to configure our model:
const MODEL_PATH = './model.json';
const INPUT_SHAPE = [28, 28, 1];
const NUM_CLASSES = 10;
The MODEL_PATH
constant defines the path to the file that will contain our trained model. The INPUT_SHAPE
constant defines the shape of the input data that our model will be trained on. The NUM_CLASSES
constant defines the number of output classes that our model will be trained to predict.
Next, we'll define a function that will build our machine learning model:
function buildModel() {
// We're building a simple machine learning model here
const model = tf.sequential();
// The first layer of our model is a convolutional layer
model.add(tf.layers.conv2d({
inputShape: INPUT_SHAPE,
kernelSize: 3,
filters: 16,
activation: 'relu'
}));
// The second layer of our model is a pooling layer
model.add(tf.layers.maxPooling2d({poolSize: 2, strides: 2}));
// The third layer of our model is a flatten layer
model.add(tf.layers.flatten());
// The fourth layer of our model is a dense layer
model.add(tf.layers.dense({units: NUM_CLASSES, activation: 'softmax'}));
// We need to compile our model before we can train it
model.compile({
optimizer: tf.train.adam(),
loss: 'categoricalCrossentropy',
metrics: ['accuracy']
});
return model;
}
This function defines a simple machine learning model consisting of four layers:
The convolutional layer is responsible for extracting features from the input data. The pooling layer is responsible for downsampling the data. The flatten layer is responsible for flattening the data. The dense layer is responsible for mapping the extracted features to the output classes.
Once we have our model defined, we need to compile it before we can train it. Compiling a model configures the model for training. We need to specify the optimizer, loss function, and metrics that we want to use.
In this case, we're using the adam
optimizer, the categoricalCrossentropy
loss function, and the accuracy
metric.
Now that we have our model defined, let's train it. We'll start by loading the training data.
The training data is stored in the ./data/train.csv
file. This file contains 28x28 images of handwritten digits, along with their labels.
We'll use the fs
module to read the contents of this file:
const trainData = fs.readFileSync('./data/train.csv', {encoding: 'utf8'});
Next, we'll parse this data into an array of tf.tensor
s:
const trainTensors = tf.tensor2d(trainData.split('\n').map(row => row.split(',').map(Number)));
We now have our training data loaded and ready to use.
Next, we'll define a function that will train our model:
async function trainModel(model, trainTensors) {
// We're training our model for 10 epochs
const epochs = 10;
// We're using a batch size of 32
const batchSize = 32;
// We're using the fit method to train our model
await model.fit(trainTensors, {
epochs: epochs,
batchSize: batchSize
});
}
This function trains our model for 10 epochs using a batch size of 32.
Once our model is trained, we'll save it to the MODEL_PATH
file:
model.save(MODEL_PATH);
Now that our model is trained, let's evaluate it. We'll start by loading the test data.
The test data is stored in the ./data/test.csv
file. This file contains 28x28 images of handwritten digits, along with their labels.
We'll use the fs
module to read the contents of this file:
const testData = fs.readFileSync('./data/test.csv', {encoding: 'utf8'});
Next, we'll parse this data into an array of tf.tensor
s:
const testTensors = tf.tensor2d(testData.split('\n').map(row => row.split(',').map(Number)));
We now have our test data loaded and ready to use.
Next, we'll define a function that will evaluate our model:
async function evaluateModel(model, testTensors) {
// We're using the evaluate method to evaluate our model
const results = await model.evaluate(testTensors);
// We're logging the loss and accuracy of our model
console.log(`loss: ${results[0]} - accuracy: ${results[1]}`);
}
This function evaluates our model on the test data.
Now that our model is trained and evaluated, let's use it to make predictions.
We'll start by loading the image that we want to make a prediction on. The image is stored in the ./data/image.png
file.
We'll use the fs
module to read the contents of this file:
const imageData = fs.readFileSync('./data/image.png');
Next, we'll parse this data into a tf.tensor
:
const imageTensor = tf.tensor3d(imageData, [28, 28, 1]);
We now have our image data loaded and ready to use.
Next, we'll define a function that will make a prediction using our model:
async function predict(model, imageTensor) {
// We're using the predict method to make a prediction
const prediction = await model.predict(imageTensor);
// We're logging the prediction that our model made
console.log(prediction);
}
This function makes a prediction on the image data.
In this post, we've seen how to build end-to-end machine learning workflows with TensorFlow.js and Node.js.
We've seen how to set up our development environment, how to create a new project, how to build a machine learning model, how to train the model, how to evaluate the model, and how to use the model to make predictions.
If you're interested in learning more about TensorFlow.js and Node.js, here are some resources that you might find helpful: