Your first Machine Learning application with Tensorflowjs and React

Start the machine learning adventure with Tensorflowjs and React, tutorial with demo and sources

Amichaud Romain

Published in

Towards Data Science

6 min readNov 27, 2020

Goals

Discover some api’s function of Tensorflowjs
Integrate Tensorflowjs in React
Create and train a basic machine learning model
Predict a value from linear data

Demo of application is available and sources are on my github.

Context

React is one of the most used javascript libraries to build a rich and powerful user interface. Tensorflowjs is the implementation of Tensorflow for NodeJS and to allow developers to use client-side machine learning models created on the server side.

Environment setup

Node.js must be installed to generate the project and manage dependencies.

The application is developped with React library, we will use Create React App to create quickly the project, with a minimal configuration.

Let’s create the project :

npx create-react-app tensorflowjs-app && cd tensorflowjs-app

Unecessary files are delete, then project look like this :

Add Tensorflowjs to the project :

npm i @tensorflow/tfjs

Write some code

Before coding, we need to define some basics use cases. User will should :

Visualize training data
Edit creation and training parameters
Chose an input value which will be use by the model to predict an outuput value

Data creation

To simplify, we will use linear data. The relation between input and output is a simple linear function which is y = 2 * x + 5. We chose this function arbitrarily.

We start from these data :

As you can see, no need machine learning to predict a value from these data.

If we want a new output value from an input x = 100, we can calculate output y = 100 * 2 + 5 = 205

The exemple is volontarly simple : the result is not important, we want to learn how we will train our model to obtain this result. Our application don’t know the relation between x (input) and y (output). It is this relation which we will create using machine learning.

Let’s create the file tensorflow.js, which will contain tensorflow functions. The first function will return our training data. We can’t work in these data directly, we need to transform them into tensors.

Tensors are objects which represents our data, they can have multiple dimensions.

// tensorflow.js
export function generateData() {
  // data may come from an external resource : CSV file, API, Database...
  const input = tf.tensor([0, 2, 4, 7, 10, 20, 50], [7, 1]);
  const label = tf.tensor([5, 9, 13, 19, 25, 45, 105], [7, 1]);
  return [input, label];
}

We create two tensors, one for x nammed “input”, and the other for y, nammed “label”.

For the user interface, we create a unique React component : TensorflowApp.js, this component will be render by our index.js. It is the equivalent of the “App.js”. I want to keep the project as simple as possible.

// TensorflowApp.js
function TensorflowApp() {
  return (
    <div>
      <h1>Tensorflowjs in React</h1>
    </div>
  );
}

Let’s get the tensors “input” and “label” :

// TensorflowApp.js
const [input, label] = generateData();

Tensorflowjs give us some functions to directly access data contained by tensors. We want to transform these tensor into arrays, we used tensor.arraySync() to do that. Then we can display our input and output like we do habitually in React.

Model creation

To create the model, we can chose some parameters, we variabilize them to let the user chose different values thanks to our React interface.

// tensorflow.js
export function createModel({
  units = 1,
  learningRate = 0.01,
  optimizer = "adam",
}) {
  const selectOptimizer = (optimizer) => {
    return OPTIMIZERS[optimizer].fn(learningRate);
  };

  const model = tf.sequential();
  model.add(tf.layers.dense({ units, inputShape: [1] }));
  model.compile({
    optimizer: selectOptimizer(optimizer),
    loss: "meanSquaredError",
  });
  return model;
}

One of the most important value is the optimizer, which is the algorithm in charge of searching the most accurate form possible for our model. It is a String, corresponding to the function’s name, we return a Tensorflow optimizer function, like a factory pattern as would have done.

To do that, we need an object containing some Tensorflowjs optimizer functions. We create this object in a different file, to let us reuse this object in other parts of our application.

// optimizers.js
export const OPTIMIZERS = {
  sgd: { libelle: "sgd", fn: (lr) => tf.train.sgd(lr) },
  adam: { libelle: "adam", fn: (lr) => tf.train.adam(lr) },
  adagrad: { libelle: "adagrad", fn: (lr) => tf.train.adagrad(lr) },
  adadelta: { libelle: "adadelta", fn: (lr) => tf.train.adadelta(lr) },
  momentum: { libelle: "momentum", fn: (lr) => tf.train.momentum(lr, 1) },
  rmsprop: { libelle: "rmsprop", fn: (lr) => tf.train.rmsprop(lr) },
};

Now we need to code the user interface, and manage user’s choices. To do that I recommand the useReducer hook provided by React.

// TensorflowApp.js
const initialModelOptions = {
  optimizer: "adam",
  learningRate: 0.01,
  epochs: 50,
};

function modelReducer(state, action) {
  switch (action.type) {
    case "SET_OPTIMIZER":
      return { ...state, optimizer: action.value };
    case "SET_LEARNINGRATE":
      return { ...state, learningRate: parseFloat(action.value) };
    case "SET_EPOCHS":
      return { ...state, epochs: parseInt(action.value) };
    default:
      return state;
  }
}

Parameters will be easily manage and updated thanks to only one function. This function will update our state with new value. This avoid to have as much “handleFunction()” as we have parameters.

// TensorflowApp.js
const handleChangeModelOptions = (action) => (e) => {
  dispatch({ type: action, value: e.target.value });
};

It only remains to code input fields corresponding to our parameters (the loss function and the learning rate).

Training our model

We must then train our model on the data, for that we create a new function. This function will take a model, an input (“x” or “input”) data as parameters, an output data (“y” or “label”), and epochs.

The epochs are the number of cycles during which the model will be trained. The larger this number is, the more refined the model and its predictions are expected to be.

// tensorflow.js
export async function trainModel(model, input, label, epochs = 150) {
  await model.fit(input, label, { epochs });
}

The model is created, it is now necessary to train it on the data.

The following function will train the model, it is an operation which may take some time and must therefore be performed asynchronously. Once the model is created, the state of the component is updated with a boolean which manages the status of our model.

// TensorflowApp.js
const handleTrain = () => {
  setModelReady(false);
  trainModel(model, input, label, modelOptions.epochs).then(() =>
    setModelReady(true)
  );
};

We still have to allow the user to set the number of learning cycles, the epochs.

Predict a value

An input field is provided to the user to allow him to enter an input value.

Results and improvments

A model is created with arbitrary parameters when launching the application. If we train this model without modifying anything and launch a prediction, the result is very unsatisfactory. If we increase the number of epochs and restart the training of the model, the predictions are immediately more correct!

To complete the application, we could display the statistics provided by Tensorflowjs on the precision of the model, or add the visualization of the data thanks to the library d3js. Also we could train our model with a web worker.

You can test the application demo, and see sources.

Originally published at https://romainamichaud.com.