15. An Introduction to TensorFlow.js – AI and Machine Learning for Coders

Chapter 15. An Introduction to TensorFlow.js

In addition to TensorFlow Lite, which enables running on native mobile or embedded systems, the TensorFlow ecosystem also includes TensorFlow.js, which lets you develop ML models using the popular JavaScript language to use directly in the browser, or on a backend with Node.js. It allows you to train new models as well as running inference on them, and includes tools that let you convert your Python-based models into JavaScript-compatible ones. In this chapter you’ll get an introduction to how TensorFlow.js fits into the overall ecosystem and a tour of its architecture, and you’ll learn how to build your own models using a free, open source IDE that integrates with your browser.

What Is TensorFlow.js?

The TensorFlow ecosystem is summarized in Figure 15-1. It comprises a suite of tools for training models, a repository for preexisting models and layers, and a set of technologies that allow you to deploy models for your end users to take advantage of.

Like TensorFlow Lite (Chapters 1214) and TensorFlow Serving (Chapter 19), TensorFlow.js mostly lives on the right side of this diagram, because while it’s primarily intended as a runtime for models, it can also be used for training models and should be considered a first-class language alongside Python and Swift for this task. TensorFlow.js can be run in the browser or on backends like Node.js, but for the purposes of this book we’ll focus primarily on the browser.

Figure 15-1. The TensorFlow ecosystem

The architecture of how TensorFlow.js gives you browser-based training and inference is shown in Figure 15-2.

Figure 15-2. TensorFlow.js high-level architecture

As a developer, you’ll typically use the Layers API, which gives Keras-like syntax in JavaScript, allowing you to use the skills you learned at the beginning of this book in JavaScript. This is underpinned by the Core API, which gives, as its name suggests, the core TensorFlow functionality in JavaScript. As well as providing the basis for the Layers API, it allows you to reuse existing Python-based models by means of a conversion toolkit that puts them into a JSON-based format for easy consumption.

The Core API then can run in a web browser, taking advantage of GPU-based acceleration using WebGL, or on Node.js, where, depending on the configuration of the environment, it can take advantage of TPU- or GPU-based acceleration in addition to the CPU.

If you’re not used to web development in either HTML or JavaScript, don’t worry; this chapter will serve as a primer, giving you enough background to help you build your first models. While you can use any web/JavaScript development environment you like, I would recommend one called Brackets to new users. In the next section you’ll see how to install that and get it up and running, after which you’ll build your first model.

Installing and Using the Brackets IDE

Brackets is a free, open source text editor that is extremely useful for web developers—in particular new ones—in that it integrates neatly with the browser, allowing you to serve your files locally so you can test and debug them. Often, when setting up web development environments, that’s the tricky part. It’s easy to write HTML or JavaScript code, but without a server to serve them to your browser, it’s hard to really test and debug them. Brackets is available for Windows, Mac, and Linux, so whatever operating system you’re using, the experience should be similar. For this chapter, I tried it out on Mint Linux, and it worked really well!

After you download and install Brackets, run it and you’ll see the Getting Started page, similar to Figure 15-3. In the top-right corner you’ll see a lightning bolt icon.

Figure 15-3. Brackets welcome screen

Click that and your web browser will launch. As you edit the HTML code in Brackets, the browser will live update. So, for example, if you change the code on line 13 that says:

<h1>GETTING STARTED WITH BRACKETS</h1>

to something else, such as:

<h1>Hello, TensorFlow Readers!</h1>

You’ll see the contents in the browser change in real time to match your edits, as shown in Figure 15-4.

Figure 15-4. Real-time updates in the browser

I find this really handy for HTML and JavaScript development in the browser, because the environment just gets out of the way and lets you focus on your code. With so many new concepts, particularly in machine learning, this is invaluable because it helps you work without too many distractions.

You’ll notice, on the Getting Started page, that you’re working in just a plain directory that Brackets serves files from. If you want to use your own, just create a directory in your filesystem and open that. New files you create in Brackets will be created and run from there. Make sure it’s a directory you have write access to so you can save your work!

Now that you have a development environment up and running, it’s time to create your first machine learning model in JavaScript. For this we’ll go back to our “Hello World” scenario where you train a model that infers the relationship between two numbers. If you’ve been working through this book from the beginning, you’ve seen this model many times already, but it’s still a useful one for helping you understand the syntactical differences you’ll need to consider when programming in JavaScript!

Building Your First TensorFlow.js Model

Before using TensorFlow.js in the browser, you’ll need to host the JavaScript in an HTML file. Create one and populate it with this skeleton HTML:

<html>
<head></head>
<body>
    <h1>First HTML Page</h1>
</body>
</html>

Then, beneath the <head> section and before the <body> tag, you can insert a <script> tag specifying the location of the TensorFlow.js library:

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"></script>

If you were to run the page now, TensorFlow.js would be downloaded, but you wouldn’t see any impact.

Next, add another <script> tag immediately below the first one. Within it, you can create a model definition. Note that while it’s very similar to how you would do it in TensorFlow in Python (refer back to Chapter 1 for details), there are some differences. For example, in JavaScript every line ends with a semicolon. Also, parameters to functions such as model.add or model.compile are in JSON notation.

This model is the familiar “Hello World” one, consisting of a single layer with a single neuron. It will be compiled with mean squared error as the loss function and stochastic gradient descent as the optimizer:

<script lang="js">        
    const model = tf.sequential();
    model.add(tf.layers.dense({units: 1, inputShape: [1]}));
    model.compile({loss:'meanSquaredError', optimizer:'sgd'});

Next, you can add the data. This is a little different from Python, where you had Numpy arrays. Of course, these aren’t available in JavaScript, so you’ll use the tf.tensor2d structure instead. It’s close, but there’s one key difference:

const xs = tf.tensor2d([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0], [6, 1]);
const ys = tf.tensor2d([-3.0, -1.0, 2.0, 3.0, 5.0, 7.0], [6, 1]);

Note that, as well as the list of values, you also have a second array that defines the shape of the first one. So, your tensor2d is initialized with a 6 × 1 list of values, followed by an array containing [6,1]. If you were to feed seven values in, the second parameter would be [7,1].

To do the training, you can then create a function called doTraining. This will train the model using model.fit, and, as before, the parameters to it will be formatted as a JSON list:

  async function doTraining(model){
    const history = 
        await model.fit(xs, ys, 
                        { epochs: 500,
                          callbacks:{
                              onEpochEnd: async(epoch, logs) =>{
                                  console.log("Epoch:" 
                                              + epoch 
                                              + " Loss:" 
                                              + logs.loss);
                                  }
                              }
                        });
}

It’s an asynchronous operation—the training will take a period of time—so it’s best to create this as an async function. Then you await the model.fit, passing it the number of epochs as a parameter. You can also specify a callback that will write out the loss for each epoch when it ends.

The last thing to do is call this doTraining method, passing it the model and reporting on the result after it finishes training:

  doTraining(model).then(() => {
    alert(model.predict(tf.tensor2d([10], [1,1])));
});

This calls model.predict, passing it a single value to get a prediction from. Because it is using a tensor2d as well as the value to predict, you also have to pass a second parameter with the shape of the first. So to predict the result for 10, you create a tensor2d with this value in an array and then pass in the shape of that array.

For convenience, here’s the complete code:

<html>
<head></head>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"></script>
<script lang="js">
    async function doTraining(model){
        const history = 
            await model.fit(xs, ys, 
                            { epochs: 500,
                              callbacks:{
                                  onEpochEnd: async(epoch, logs) =>{
                                      console.log("Epoch:" 
                                                  + epoch 
                                                  + " Loss:" 
                                                  + logs.loss);
                                  
                                  }
                              }
                            });
    }
    const model = tf.sequential();
    model.add(tf.layers.dense({units: 1, inputShape: [1]}));
    model.compile({loss:'meanSquaredError', 
                   optimizer:'sgd'});
    model.summary();
    const xs = tf.tensor2d([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0], [6, 1]);
    const ys = tf.tensor2d([-3.0, -1.0, 2.0, 3.0, 5.0, 7.0], [6, 1]);
    doTraining(model).then(() => {
        alert(model.predict(tf.tensor2d([10], [1,1])));
    });
</script>
<body>
    <h1>First HTML Page</h1>
</body>
</html>

When you run this page, it will appear as if nothing has happened. Wait a few seconds, and then a dialog will appear like the one in Figure 15-5. This is the alert dialog that is shown with the results of the prediction for [10].

Figure 15-5. Results of the inference after training

It can be a little disconcerting to have that long pause before the dialog is shown—as you’ve probably guessed, the model was training during that period. Recall that in the doTraining function, you created a callback that writes the loss per epoch to the console log. If you want to see this, you can do so with the browser’s developer tools. In Chrome you can access these by clicking the three dots in the upper-right corner and selecting More Tools → Developer Tools, or by pressing Ctrl-Shift-I.

Once you have them, select Console at the top of the pane and refresh the page. As the model retrains, you’ll see the loss per epoch (see Figure 15-6).

Figure 15-6. Exploring the per-epoch loss in the browser’s developer tools

Now that you’ve gone through your first (and simplest) model, you’re ready to build something a little more complex.

Creating an Iris Classifier

The last example was a very simple one, so let’s work on one that’s a little more complex next. If you’ve done any work with machine learning, you’ve probably heard about the Iris dataset, which is a perfect one for learning ML.

The dataset contains 150 data items with four attributes each, describing three classes of flower. The attributes are sepal length and width, and petal length and width. When plotted against each other, clear clusters of flower types are seen (Figure 15-7).

Figure 15-7. Plotting features in the Iris dataset (source: Nicoguaro, available on Wikimedia Commons)

Figure 15-7 shows the complexity of the problem, even with a simple dataset like this. How does one separate the three types of flowers using rules? The petal length versus petal width plots come close, with the Iris setosa samples (in red) being very distinct from the others, but the blue and green sets are intertwined. This makes for an ideal learning set in ML: it’s small, so fast to train, and you can use it to solve a problem that’s difficult to do in rules-based programming!

You can download the dataset from the UCI Machine Learning Repository, or use the version in the book’s GitHub repository, which I’ve converted to CSV to make it easier to use in JavaScript.

The CSV looks like this:

sepal_length,sepal_width,petal_length,petal_width,species
5.1,3.5,1.4,0.2,setosa
4.9,3,1.4,0.2,setosa
4.7,3.2,1.3,0.2,setosa
4.6,3.1,1.5,0.2,setosa
5,3.6,1.4,0.2,setosa
5.4,3.9,1.7,0.4,setosa
4.6,3.4,1.4,0.3,setosa
5,3.4,1.5,0.2,setosa
...

The four data points for each flower are the first four values. The label is the fifth value, one of setosa, versicolor, or virginica. The first line in the CSV file contains the column labels. Keep that in mind—it will be useful later!

To get started, create a basic HTML page as before, and add the <script> tag to load TensorFlow.js:

<html>
<head></head>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"></script>
<body>
    <h1>Iris Classifier</h1>
</body>
</html>

To load a CSV file, TensorFlow.js has the tf.data.csv API that you can give a URL to. This also allows you to specify which column is the label. Because the first line in the CSV file I’ve prepared contains column names, you can specify which one contains the labels, which in this case is species, as follows:

<script lang="js">
    async function run(){
        const csvUrl = 'iris.csv';
        const trainingData = tf.data.csv(csvUrl, {
            columnConfigs: {
                species: {
                    isLabel: true
                }
            }
        });

The labels are strings, which you don’t really want to train a neural network with. This is going to be a multiclass classifier with three output neurons, each containing the probability that the input data represents the respective species of flower. Thus, a one-hot encoding of the labels is perfect.

This way, if you represent setosa as [1, 0, 0], showing that you want the first neuron to light up for this class, virginica as [0, 1, 0], and versicolor as [0, 0, 1], you’re effectively defining the template for how the final layer neurons should behave for each class.

Because you used tf.data to load the data, you have the ability to use a mapping function that will take in your xs (features) and ys (labels) and map them differently. So, to keep the features intact and one-hot encode the labels, you can write code like this:

  const convertedData =
    trainingData.map(({xs, ys}) => {
        const labels = [
            ys.species == "setosa" ? 1 : 0,
            ys.species == "virginica" ? 1 : 0,
            ys.species == "versicolor" ? 1 : 0
        ] 
        return{ xs: Object.values(xs), ys: Object.values(labels)};
    }).batch(10);

Note that the labels are stored as an array of three values. Each value defaults to 0, unless the species matches the given string, in which case it will be a 1. Thus, setosa will be encoded as [1, 0, 0], and so on.

The mapping function will return the xs unaltered and the ys one-hot encoded.

You can now define your model. Your input layer’s shape is the number of features, which is the number of columns in the CSV file minus 1 (because one of the columns represents the labels):

const numOfFeatures = (await trainingData.columnNames()).length - 1;

const model = tf.sequential();
model.add(tf.layers.dense({inputShape: [numOfFeatures], 
                           activation: "sigmoid", units: 5}))

model.add(tf.layers.dense({activation: "softmax", units: 3}));

Your final layer has three units because of the three classes that are one-hot encoded in the training data.

Next, you’ll specify the loss function and optimizer. Because this is a multiple-category classifier, make sure you use a categorical loss function such as categorical cross entropy. You can use an optimizer such as adam in the tf.train namespace and pass it parameters like the learning rate (here, 0.06):

model.compile({loss: "categoricalCrossentropy", 
    optimizer: tf.train.adam(0.06)});

Because the data was formatted as a dataset, you can use model.fitDataset for training instead of model.fit. To train for one hundred epochs and catch the loss in the console, you can use a callback like this:

await model.fitDataset(convertedData, 
                       {epochs:100,
                        callbacks:{
                            onEpochEnd: async(epoch, logs) =>{
                                console.log("Epoch: " + epoch + 
                                 " Loss: " + logs.loss);
                            }
                        }});

To test the model after it has finished training, you can load values into a tensor2d. Don’t forget that when using a tensor2d, you also have to specify the shape of the data. In this case, to test a set of four values, you would define them in a tensor2d like this:

const testVal = tf.tensor2d([4.4, 2.9, 1.4, 0.2], [1, 4]);

You can then get a prediction by passing this to model.predict:

const prediction = model.predict(testVal);

You’ll get a tensor value back that looks something like this:

[[0.9968228, 0.00000029, 0.0031742],]

To get the largest value, you can use the argMax function:

tf.argMax(prediction, axis=1)

This will return [0] for the preceding data because the neuron at position 0 had the highest probability.

To unpack this into a value, you can use .dataSync. This operation synchronously downloads a value from the tensor. It does block the UI thread, so be careful when using it!

The following code will simply return 0 instead of [0]:

const pIndex = tf.argMax(prediction, axis=1).dataSync();

To map this back to a string with the class name, you can then use this code:

const classNames = ["Setosa", "Virginica", "Versicolor"];
alert(classNames[pIndex])

Now you’ve seen how to load data from a CSV file, turn it into a dataset, and then fit a model from that dataset, as well as how to run predictions from that model. You’re well equipped to experiment with other datasets of your choosing to further hone your skills!

Summary

This chapter introduced you to TensorFlow.js and how it can be used to train models and perform inference in the browser. You saw how to use the open source Brackets IDE to code and test your models on a local web server, and used this to train your first two models: a “Hello World” linear regression model and a basic classifier for the popular Iris dataset. These were very simple scenarios, but in Chapter 16 you’ll take things to the next level and see how to train computer vision models with TensorFlow.js.