Set up a TensorFlow.js project

This document shows you how to install and use TensorFlow.js in a browser environment and in Node.js.

Browser setup

There are two recommended ways to use TensorFlow.js in a browser-based project:

If you're new to web development, or you haven't used JavaScript build tools before, you might want to try the script tag approach first. If you usually bundle or process your web assets, or you plan to write larger applications, you should consider using build tools.

Use a script tag

To get TensorFlow.js using a script tag, add the following to your main HTML file:

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest/dist/tf.min.js"></script>

The following example shows how to define and train a model in the browser:

<!DOCTYPE html>
<html lang="en-US">
  <head>
    <meta charset="utf-8" />
    <title>TensorFlow.js browser example</title>

    <!-- Load TensorFlow.js from a script tag -->
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest/dist/tf.min.js"></script>
  </head>
  <body>
    <h1>TensorFlow.js example</h1>
    <h2>Open the console to see the results.</h2>
    <script>
    // Define a model for linear regression. The script tag makes `tf` available
    // as a global variable.
    const model = tf.sequential();
    model.add(tf.layers.dense({units: 1, inputShape: [1]}));

    model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});

    // Generate some synthetic data for training.
    const xs = tf.tensor2d([1, 2, 3, 4], [4, 1]);
    const ys = tf.tensor2d([1, 3, 5, 7], [4, 1]);

    // Train the model using the data.
    model.fit(xs, ys, {epochs: 10}).then(() => {
      // Use the model to do inference on a data point the model hasn't seen before:
      model.predict(tf.tensor2d([5], [1, 1])).print();
      // Open the browser devtools to see the output
    });
    </script>
  </body>
</html>

To run the example, follow these steps:

  1. Save the example document in a file called index.html.

  2. Double-click index.html to open it in your default browser.

    Alternatively, you can serve index.html by running npx http-server in the same directory as index.html. (If you're prompted for permission to install http-server, enter y.) Then go to http://localhost:8080 in your browser.

  3. Open the browser console to see the output of the script.

  4. Refresh the page to see a new (and very probably different) prediction.

Install from NPM

To install TensorFlow.js from NPM, use either the npm CLI or yarn.

NPM

npm install @tensorflow/tfjs

Yarn

yarn add @tensorflow/tfjs

The following example shows how to import TensorFlow.js, define a model, and train the model.

import * as tf from '@tensorflow/tfjs';

// Define a model for linear regression.
const model = tf.sequential();
model.add(tf.layers.dense({units: 1, inputShape: [1]}));

model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});

// Generate some synthetic data for training.
const xs = tf.tensor2d([1, 2, 3, 4], [4, 1]);
const ys = tf.tensor2d([1, 3, 5, 7], [4, 1]);

// Train the model using the data.
model.fit(xs, ys, {epochs: 10}).then(() => {
  // Use the model to do inference on a data point the model hasn't seen before:
  model.predict(tf.tensor2d([5], [1, 1])).print();
  // Open the browser devtools to see the output
});

Node.js setup

To use TensorFlow.js in Node.js, use either the npm CLI or yarn to complete one of the installation options below.

To learn more about using TensorFlow.js in Node.js, see the Node.js guide. For additional installation information, see the TensorFlow.js for Node.js repository.

Option 1: Install TensorFlow.js with native C++ bindings.

The tfjs-node module provides native TensorFlow execution in JavaScript applications under the Node.js runtime, accelerated by the TensorFlow C binary.

Install tfjs-node:

NPM

npm install @tensorflow/tfjs-node

Yarn

yarn add @tensorflow/tfjs-node

The following example shows how to import tfjs-node, define a model, and train the model.

// Use `tfjs-node`. Note that `tfjs` is imported indirectly by `tfjs-node`.
const tf = require('@tensorflow/tfjs-node');

// Define a simple model.
const model = tf.sequential();
model.add(tf.layers.dense({units: 100, activation: 'relu', inputShape: [10]}));
model.add(tf.layers.dense({units: 1, activation: 'linear'}));
model.compile({optimizer: 'sgd', loss: 'meanSquaredError'});

const xs = tf.randomNormal([100, 10]);
const ys = tf.randomNormal([100, 1]);

// Train the model.
model.fit(xs, ys, {
  epochs: 100,
  callbacks: {
    onEpochEnd: (epoch, log) => console.log(`Epoch ${epoch}: loss = ${log.loss}`)
  }
});

Option 2: Install TensorFlow.js for GPU

(Linux only) If your system has an NVIDIA® GPU with CUDA support, you can use the GPU package for improved performance.

Install tfjs-node-gpu:

NPM

npm install @tensorflow/tfjs-node-gpu

Yarn

yarn add @tensorflow/tfjs-node-gpu

The following example shows how to import tfjs-node-gpu, define a model, and train the model.

// Use `tfjs-node-gpu`. Note that `tfjs` is imported indirectly by `tfjs-node-gpu`.
const tf = require('@tensorflow/tfjs-node-gpu');

// Define a simple model.
const model = tf.sequential();
model.add(tf.layers.dense({units: 100, activation: 'relu', inputShape: [10]}));
model.add(tf.layers.dense({units: 1, activation: 'linear'}));
model.compile({optimizer: 'sgd', loss: 'meanSquaredError'});

const xs = tf.randomNormal([100, 10]);
const ys = tf.randomNormal([100, 1]);

// Train the model.
model.fit(xs, ys, {
  epochs: 100,
  callbacks: {
    onEpochEnd: (epoch, log) => console.log(`Epoch ${epoch}: loss = ${log.loss}`)
  }
});

Option 3: Install the pure JavaScript version

The tfjs module is the same package that you'd use in the browser. It's the slowest of the Node.js options in terms of performance.

Install tfjs:

NPM

npm install @tensorflow/tfjs

Yarn

yarn add @tensorflow/tfjs

The following example shows how to import tfjs, define a model, and train the model.

// Use `tfjs`.
const tf = require('@tensorflow/tfjs');

// Define a simple model.
const model = tf.sequential();
model.add(tf.layers.dense({units: 100, activation: 'relu', inputShape: [10]}));
model.add(tf.layers.dense({units: 1, activation: 'linear'}));
model.compile({optimizer: 'sgd', loss: 'meanSquaredError'});

const xs = tf.randomNormal([100, 10]);
const ys = tf.randomNormal([100, 1]);

// Train the model.
model.fit(xs, ys, {
  epochs: 100,
  callbacks: {
    onEpochEnd: (epoch, log) => console.log(`Epoch ${epoch}: loss = ${log.loss}`)
  }
});

TypeScript

If you're using TensorFlow.js in a TypeScript project, and you have strict null checking enabled, you might need to set skipLibCheck: true in your tsconfig.json to avoid errors during compilation.