Use a pre-trained model  |  TensorFlow.js
Skip to main content
Install
Learn
Introduction
New to TensorFlow?
Tutorials
Learn how to use TensorFlow with end-to-end examples
Guide
Learn framework concepts and components
Learn ML
Educational resources to master your path with TensorFlow
API
TensorFlow (v2.16.1)
Versions…
TensorFlow.js
TensorFlow Lite
TFX
Ecosystem
LIBRARIES
TensorFlow.js
Develop web ML applications in JavaScript
TensorFlow Lite
Deploy ML on mobile, microcontrollers and other edge devices
TFX
Build production ML pipelines
All libraries
Create advanced models and extend TensorFlow
RESOURCES
Models & datasets
Pre-trained models and datasets built by Google and the community
Tools
Tools to support and accelerate TensorFlow workflows
Responsible AI
Resources for every stage of the ML workflow
Recommendation systems
Build recommendation systems with open source tools
Community
Groups
User groups, interest groups and mailing lists
Contribute
Guide for contributing to code and documentation
Blog
Stay up to date with all things TensorFlow
Forum
Discussion platform for the TensorFlow community
Why TensorFlow
About
Case studies
/
English
Español
Español – América Latina
Français
Indonesia
Italiano
Polski
Português
Português – Brasil
Tiếng Việt
Türkçe
Русский
עברית
العربيّة
فارسی
हिंदी
বাংলা
ภาษาไทย
中文 – 简体
日本語
한국어
GitHub
Sign in
For JavaScript
Overview
Tutorials
Guide
Models
Demos
API
Install
Learn
More
Overview
Tutorials
Guide
Models
Demos
API
API
More
Ecosystem
More
Community
More
Why TensorFlow
More
GitHub
Get started
Set up a TensorFlow.js project
Upgrade to TensorFlow.js 3.0
Train models
Fit a curve to two-dimensional data
Recognize handwritten digits with CNNs
Predict baseball pitch types in Node.js
Train a model using a web worker
Transfer learning
What is transfer learning?
Build an image classifier
Build an audio recognizer
Import Python models
Use a pre-trained model
Import a Keras model
Import a TensorFlow SavedModel
Deployment & Optimization
Generating size-optimized browser bundles
Deploy a Chrome extension
Web ML applications
Predictive prefetching
React Native
Introduction
Tutorials
Guide
Learn ML
TensorFlow (v2.16.1)
Versions…
TensorFlow.js
TensorFlow Lite
TFX
LIBRARIES
TensorFlow.js
TensorFlow Lite
TFX
All libraries
RESOURCES
Models & datasets
Tools
Responsible AI
Recommendation systems
Groups
Contribute
Blog
Forum
About
Case studies
New to machine learning? Watch a video course to get practical working knowledge of ML using web technologiesView series
TensorFlow
Learn
For JavaScript
Tutorials
Use a pre-trained model
Stay organized with collections
Save and categorize content based on your preferences.
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. 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,
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. If you don't want to set up a
development environment, you can skip to
Explore the example.
The example code is available on
GitHub.
(Optional) Run the example locally
Prerequisites
To run the example app locally, you need the following installed in your
development environment:
Node.js (download)
Yarn (install)
Install and run the example app
Clone or download the
tfjs-examples repository.
Change into the mnist-transfer-cnn directory:
cd tfjs-examples/mnist-transfer-cnn
Install dependencies:
yarn
Start the development server:
yarn run watch
Explore the example
Open the example app.
(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.
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.
The
index.js
file contains most of the training code for the demo. When index.js runs in
the browser, a setup function,
setupMnistTransferCNN,
instantiates and initializes
MnistTransferCNNPredictor,
which encapsulates the retraining and prediction routines.
The initialization method,
MnistTransferCNNPredictor.init,
loads a model, loads retraining data, and creates test data. Here's
the line
that loads the model:
this.model = await loader.loadHostedPretrainedModel(urls.model);
If you look at the definition of
loader.loadHostedPretrainedModel,
you'll see that it returns the result of a call to
tf.loadLayersModel.
This is the TensorFlow.js API for loading a model composed of Layer objects.
The retraining logic is defined in
MnistTransferCNNPredictor.retrainModel.
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,
and then it's
trained
on the test data using model.fit():
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.
After being trained on the new dataset (digits 5-9), the model can be used to
make predictions. The
MnistTransferCNNPredictor.predict
method does this using
model.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, 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
Import a TensorFlow model into TensorFlow.js
Pre-made models for TensorFlow.js
TensorFlow Core
Keras: Transfer learning and fine-tuning
Transfer learning and fine-tuning
Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates.
Last updated 2023-05-03 UTC.
[[["Easy to understand","easyToUnderstand","thumb-up"],["Solved my problem","solvedMyProblem","thumb-up"],["Other","otherUp","thumb-up"]],[["Missing the information I need","missingTheInformationINeed","thumb-down"],["Too complicated / too many steps","tooComplicatedTooManySteps","thumb-down"],["Out of date","outOfDate","thumb-down"],["Samples / code issue","samplesCodeIssue","thumb-down"],["Other","otherDown","thumb-down"]],["Last updated 2023-05-03 UTC."],[],[]]
Stay connected
Blog
Forum
GitHub
Twitter
YouTube
Support
Issue tracker
Release notes
Stack Overflow
Brand guidelines
Cite TensorFlow
Terms
Privacy
Manage cookies
Sign up for the TensorFlow newsletter
Subscribe
English
Español
Español – América Latina
Français
Indonesia
Italiano
Polski
Português
Português – Brasil
Tiếng Việt
Türkçe
Русский
עברית
العربيّة
فارسی
हिंदी
বাংলা
ภาษาไทย
中文 – 简体
日本語
한국어