Generating size-optimized browser bundles with 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
Generating size-optimized browser bundles with TensorFlow.js
Stay organized with collections
Save and categorize content based on your preferences.
Overview
TensorFlow.js 3.0 brings support for building size-optimized, production oriented browser bundles. To put it another way we want to make it easier for you to ship less JavaScript to the browser.
This feature is geared towards users with production use cases who would particularly benefit from shaving bytes off their payload (and are thus willing to put in the effort to achieve this). To use this feature you should be familiar with ES Modules, JavaScript bundling tools such as webpack or rollup, and concepts such as tree-shaking/dead-code elimination.
This tutorial demonstrates how to create a custom tensorflow.js module that can be used with a bundler to generate a size optimized build for a program using tensorflow.js.
Terminology
In the context of this document there are a few key terms we will be using:
ES Modules - The standard JavaScript module system. Introduced in ES6/ES2015. Identifiable by use of import and export statements.
Bundling - Taking a set of JavaScript assets and grouping/bundling them into one or more JavaScript assets that are usable in a browser. This is the step that usually produces the final assets that are served to the browser. Applications will generally do their own bundling directly from transpiled library sources. Common bundlers include rollup and webpack. The end result of bundling is a known as a bundle (or sometimes as a chunk if it is split into multiple parts)
Tree-Shaking / Dead Code Elimination - Removal of code that is not used by the final written application. This is done during bundling, typically in the minification step.
Operations (Ops) - A mathematical operation on one or more tensors that produces one or more tensors as output. Ops are ‘high level’ code and can use other ops to define their logic.
Kernel - A specific implementation of an op tied to specific hardware capabilities. Kernels are ‘low level’  and backend specific. Some ops have a one-to-one mapping from op to kernel while other ops use multiple kernels.
Scope and Use Cases
Inference only graph-models
The primary use case we heard about from users related to this, and are supporting in this release is that of doing inference with TensorFlow.js graph models. If you are using a TensorFlow.js layers model, you can convert this to the graph-model format using the tfjs-converter. The graph model format is more efficient for the inference use case.
Low level Tensor manipulation with tfjs-core
The other use case we support is programs that directly use the @tensorflow/tjfs-core package for lower level tensor manipulation.
Our approach to custom builds
Our core principles when designing this functionality includes the following:
Make maximal use of the JavaScript module system (ESM) and allow users of TensorFlow.js to do the same.
Make TensorFlow.js as tree-shakeable as possible by existing bundlers (e.g. webpack, rollup, etc). This enables users to take advantage of all the capabilities of those bundlers including features like code splitting.
As much as possible maintain ease of use for users who are not as sensitive to bundle size. This does mean that production builds will require more effort as many of the defaults in our libraries support ease of use over size optimized builds.
The primary goal of our workflow is to produce a custom JavaScript module for TensorFlow.js that contains only the functionality required for the program we are trying to optimize. We rely on existing bundlers to do the actual optimization.
While we primarily rely on the JavaScript module system, we also provide a custom CLI tool to handle parts that aren’t easy to specify through the module system in user facing code. Two examples of this are:
Model specifications stored in model.json files
The op to backend-specific-kernel dispatching system we use.
This makes generating a custom tfjs build a bit more involved than just pointing a bundler to the regular @tensorflow/tfjs package.
How to create size optimized custom bundles
Step 1: Determine which kernels your program is using
This step lets us determine all the kernels used by any models you run or pre/post-processing code given the backend you have selected.
Use tf.profile to run the parts of your application that use tensorflow.js and get the kernels. It will look something like this
const profileInfo = await tf.profile(() => {
// You must profile all uses of tf symbols.
runAllMyTfjsCode();
});
const kernelNames = profileInfo.kernelNames
console.log(kernelNames);
Copy that list of kernels to your clipboard for the next step.
You need to profile the code using the same backend(s) that you want to use in your custom bundle.
You will need to repeat this step if your model changes or your pre/post-processing code changes.
Step 2. Write a configuration file for the custom tfjs module
Here is an example config file.
It looks like this:
{
"kernels": ["Reshape", "_FusedMatMul", "Identity"],
"backends": [
"cpu"
],
"models": [
"./model/model.json"
],
"outputPath": "./custom_tfjs",
"forwardModeOnly": true
}
kernels: The list of kernels to include in the bundle. Copy this from the output of Step 1.
backends: The list of backend(s) you want to include. Valid options include "cpu", "webgl" and “wasm”.
models: A list of model.json files for models you load in your application. Can be empty if your program does not use tfjs_converter to load a graph model.
outputPath: A path to a folder to put the generated modules in.
forwardModeOnly: Set this to false if you want to include gradients for the kernels listed prior.
Step 3. Generate the custom tfjs module
Run the custom build tool with the config file as an argument. You need to have the @tensorflow/tfjs package installed in order to have access to this tool.
npx tfjs-custom-module  --config custom_tfjs_config.json
This will create a folder at outputPath with some new files.
Step 4. Configure your bundler to alias tfjs to the new custom module.
In bundlers like webpack and rollup we can alias the existing references to tfjs modules to point to our newly generated custom tfjs modules. There are three modules that need to be aliased for maximum savings in bundle size.
Here is an snippet of what that looks like in webpack (full example here):
...
config.resolve = {
alias: {
'@tensorflow/tfjs$':
path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
'@tensorflow/tfjs-core$': path.resolve(
__dirname, './custom_tfjs/custom_tfjs_core.js'),
'@tensorflow/tfjs-core/dist/ops/ops_for_converter': path.resolve(
__dirname, './custom_tfjs/custom_ops_for_converter.js'),
}
}
...
And here is the equivalent code snippet for rollup (full example here):
import alias from '@rollup/plugin-alias';
...
alias({
entries: [
{
find: /@tensorflow\/tfjs$/,
replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
},
{
find: /@tensorflow\/tfjs-core$/,
replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs_core.js'),
},
{
find: '@tensorflow/tfjs-core/dist/ops/ops_for_converter',
replacement: path.resolve(__dirname, './custom_tfjs/custom_ops_for_converter.js'),
},
],
}));
...
If your bundler does not support module aliasing, you will need to change your import statements to import tensorflow.js from generated custom_tfjs.js that was created in Step 3. Op definitions will not be tree-shaken out, but kernels still will be tree-shaken. Generally tree-shaking kernels is what provides the biggest savings in final bundle size.
If you are only using the @tensoflow/tfjs-core package, then you only need to alias that one package.
Step 5. Create your bundle
Run your bundler (e.g. webpack or rollup) to produce your bundle. The size of the bundle should be smaller than if you run the bundler without module aliasing. You can also use visualizers like this one to see what made it into your final bundle.
Step 6. Test your app
Make sure to test that your app is working as expected!
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 2022-11-01 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 2022-11-01 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
Русский
עברית
العربيّة
فارسی
हिंदी
বাংলা
ภาษาไทย
中文 – 简体
日本語
한국어