# TensorFlow.js Part 1 - Develop Model in Python

Update 23 October 2022: The model should be called an autoencoder, see discussion.

This post is part one of a three-part series in which we are going to explore how TensorFlow models can be integrated in website front-ends written in JavaScript with TensorFlow.js. This is useful if you want to run a small model, but don’t want to send client data to a server out of privacy concerns. TensorFlow.js can also be used in JavaScript on a backend which runs Node.js.

In this first part, we are going to develop and train a model with TensorFlow in Python. In part two we then export this model so that it can be used in TensorFlow.js. The last part explores how a converted model can be used in JavaScript. We develop a simple website with React.js which allows us to upload an image and run the model on it.

To follow along, you can download code here. If you want to dive into the Javascript part directly, you can also skip this part and use a model from the models folder.

Additionally, the content of this post is also available as a video:

## Develop Model and Loss Function

We are going to build a simple fully convolutional network that will be trained with a pixel loss to reconstruct an input image. Such a model is used for example for real-time style transfer or image upscaling (super-resolution).

Next, we define a pixel loss which takes the squared distance between two images based on their pixel values.

## Data

To train the model, we use a few images which can be found in the data folder. For this demo, you can use as many images as you like since we are not interested in making a good model, but rather want to explore how to use models in TensorFlow.js. I used four images from the COCO dataset. The model is going to overfit on those, and that is totally fine for this tutorial.

To load data from the folder we define an image_dataset.

## Training Loop

To train the model we need one more piece: a training loop which takes images, runs a forward pass, evaluates the output with the loss function, and updates the weights with the gradients.

The function decorator @tf.function compiles the function into a static graph which means for us that it runs faster.

With the function above, we can finally define the training loop.

After running the loop, the model is saved to the models folder. With the command model.save() we save the model in the SavedModel format. An alternative is the Keras HDF5 format which requires different commands later.

## Check model

Finally, we can check if the model we just trained can reconstruct the images.

Figure 1 shows the result from the image which is in the data folder. As you can see it looks pretty good! This is no surprise since the model memorized this image due to lack of data, but it will do for this demo.

Tags:

Updated: