# Use a pre-trained model

<br />

In this tutorial you'll explore an example web application that demonstrates
transfer learning using the TensorFlow.js Layers API. The example loads a
pre-trained model and then retrains the model in the browser.

The model has been pre-trained in Python on digits 0-4 of the
[MNIST digits classification dataset](https://keras.io/api/datasets/mnist/). The
retraining (or transfer learning) in the browser uses digits 5-9. The example
shows that the first several layers of a pre-trained model can be used to
extract features from new data during transfer learning, thus enabling faster
training on the new data.

The example application for this tutorial is
[available online](https://storage.googleapis.com/tfjs-examples/mnist-transfer-cnn/dist/index.html),
so you don't need to download any code or set up a development environment. If
you'd like to run the code locally, complete the optional steps in
[Run the example locally](https://www.tensorflow.org/js/tutorials/conversion/pretrained_model#run_example_locally). If you don't want to set up a
development environment, you can skip to
[Explore the example](https://www.tensorflow.org/js/tutorials/conversion/pretrained_model#explore_example).

The example code is available on
[GitHub](https://github.com/tensorflow/tfjs-examples/tree/master/mnist-transfer-cnn).

## (Optional) Run the example locally

### Prerequisites

To run the example app locally, you need the following installed in your
development environment:

- Node.js ([download](https://nodejs.org/en/download/))
- Yarn ([install](https://yarnpkg.com/en/docs/install))

### Install and run the example app

1. Clone or download the [`tfjs-examples`](https://github.com/tensorflow/tfjs-examples) repository.
2. Change into the `mnist-transfer-cnn` directory:

       cd tfjs-examples/mnist-transfer-cnn

3. Install dependencies:

       yarn

4. Start the development server:

       yarn run watch

## Explore the example

[Open the example app](https://storage.googleapis.com/tfjs-examples/mnist-transfer-cnn/dist/index.html).
(Or, if you're running the example locally, go to `http://localhost:1234` in
your browser.)

You should see a page titled **MNIST CNN Transfer Learning**. Follow the
instructions to try the app.

Here are a few things to try:

- Experiment with the different training modes and compare loss and accuracy.
- Select different bitmap examples and inspect the classification probabilities. Note that the numbers in each bitmap example are grayscale integer values representing pixels from an image.
- Edit the bitmap integer values and see how the changes affect classification probabilities.

## Explore the code

The example web app loads a model that has been pre-trained on a subset of the
MNIST dataset. The pre-training is defined in a Python program:
[`mnist_transfer_cnn.py`](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/python/mnist_transfer_cnn.py).
The Python program is out-of-scope for this tutorial, but it's worth looking at
if you'd like to see an example of [model conversion](https://www.tensorflow.org/js/guide/conversion).

The
[`index.js`](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/index.js)
file contains most of the training code for the demo. When `index.js` runs in
the browser, a setup function,
[`setupMnistTransferCNN`](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/index.js#L168-L194),
instantiates and initializes
[`MnistTransferCNNPredictor`](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/index.js#L40-L162),
which encapsulates the retraining and prediction routines.

The initialization method,
[`MnistTransferCNNPredictor.init`](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/index.js#L44-L59),
loads a model, loads retraining data, and creates test data. Here's
[the line](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/index.js#L46)
that loads the model:

    this.model = await loader.loadHostedPretrainedModel(urls.model);

If you look at the definition of
[`loader.loadHostedPretrainedModel`](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/loader.js#L35-L54),
you'll see that it returns the result of a call to
[`tf.loadLayersModel`](https://js.tensorflow.org/api/latest/#loadLayersModel).
This is the TensorFlow.js API for loading a model composed of Layer objects.

The retraining logic is defined in
[`MnistTransferCNNPredictor.retrainModel`](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/index.js#L111-L162).
If the user has selected **Freeze feature layers** as the training mode, the
first 7 layers of the base model are frozen, and only the final 5 layers are
trained on new data. If the user has selected **Reinitialize weights**, all the
weights are reset, and the app effectively trains the model from scratch.

    if (trainingMode === 'freeze-feature-layers') {
      console.log('Freezing feature layers of the model.');
      for (let i = 0; i < 7; ++i) {
        this.model.layers[i].trainable = false;
      }
    } else if (trainingMode === 'reinitialize-weights') {
      // Make a model with the same topology as before, but with re-initialized
      // weight values.
      const returnString = false;
      this.model = await tf.models.modelFromJSON({
        modelTopology: this.model.toJSON(null, returnString)
      });
    }

The model is then
[compiled](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/index.js#L131-L135),
and then it's
[trained](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/index.js#L146-L159)
on the test data using [`model.fit()`](https://www.tensorflow.org/js/guide/train_models#modelfit):

    await this.model.fit(this.gte5TrainData.x, this.gte5TrainData.y, {
      batchSize: batchSize,
      epochs: epochs,
      validationData: [this.gte5TestData.x, this.gte5TestData.y],
      callbacks: [
        ui.getProgressBarCallbackConfig(epochs),
        tfVis.show.fitCallbacks(surfaceInfo, ['val_loss', 'val_acc'], {
          zoomToFit: true,
          zoomToFitAccuracy: true,
          height: 200,
          callbacks: ['onEpochEnd'],
        }),
      ]
    });

To learn more about the `model.fit()` parameters, see the
[API documentation](https://js.tensorflow.org/api/latest/#tf.LayersModel.fit).

After being trained on the new dataset (digits 5-9), the model can be used to
make predictions. The
[`MnistTransferCNNPredictor.predict`](https://github.com/tensorflow/tfjs-examples/blob/master/mnist-transfer-cnn/index.js#L96-L109)
method does this using
[`model.predict()`](https://js.tensorflow.org/api/latest/#tf.LayersModel.predict):

    // Perform prediction on the input image using the loaded model.
    predict(imageText) {
      tf.tidy(() => {
        try {
          const image = util.textToImageArray(imageText, this.imageSize);
          const predictOut = this.model.predict(image);
          const winner = predictOut.argMax(1);

          ui.setPredictResults(predictOut.dataSync(), winner.dataSync()[0] + 5);
        } catch (e) {
          ui.setPredictError(e.message);
        }
      });
    }

Note the use of [`tf.tidy`](https://js.tensorflow.org/api/latest/#tidy), which
helps prevent memory leaks.

## Learn more

This tutorial has explored an example app that performs transfer learning in the
browser using TensorFlow.js. Check out the resources below to learn more about
pre-trained models and transfer learning.

**TensorFlow.js**

- [Importing a Keras model into TensorFlow.js](https://www.tensorflow.org/js/tutorials/conversion/import_keras)
- [Import a TensorFlow model into TensorFlow.js](https://www.tensorflow.org/js/tutorials/conversion/import_saved_model)
- [Pre-made models for TensorFlow.js](https://www.tensorflow.org/js/guide/premade_models)

**TensorFlow Core**

- [Keras: Transfer learning and fine-tuning](https://www.tensorflow.org/guide/keras/transfer_learning)
- [Transfer learning and fine-tuning](https://www.tensorflow.org/tutorials/images/transfer_learning)