New to machine learning? Watch a video course to get practical working knowledge of ML using web technologiesView series

TensorFlow.js in Node.js

This guide describes the TensorFlow.js packages and APIs available for Node.js.

To learn how to install TensorFlow.js in Node.js, see the setup tutorial. For additional information on installation and support, see the TensorFlow.js for Node.js repository.

TensorFlow CPU

The TensorFlow CPU package can be imported as follows:

import * as tf from '@tensorflow/tfjs-node'

When you import TensorFlow.js from this package, you get a module that's accelerated by the TensorFlow C binary and runs on the CPU. TensorFlow on the CPU uses hardware acceleration to optimize linear algebra computation.

This package works on Linux, Windows, and macOS platforms where TensorFlow is supported.

Note: You do not have to import @tensorflow/tfjs or add it to your package.json. This is indirectly imported by the tfjs-node library.

TensorFlow GPU

The TensorFlow GPU package can be imported as follows:

import * as tf from '@tensorflow/tfjs-node-gpu'

Like the CPU package, the module is accelerated by the TensorFlow C binary. But the GPU package runs tensor operations on the GPU with CUDA, so it's only available on Linux. This binding can be at least an order of magnitude faster than the other binding options.

Note: This package currently only works with CUDA. You need to have CUDA installed on your machine with an NVIDIA graphics card before using tfjs-node-gpu.

Note: You do not have to import '@tensorflow/tfjs' or add it to your package.json. This is indirectly imported by the tfjs-node-gpu library.

TensorFlow for pure JavaScript

There's also a version of TensorFlow.js that runs pure JavaScript on the CPU. It can be imported as follows:

import * as tf from '@tensorflow/tfjs'

This package is the same package that you'd use in the browser. In this package, the operations are run in vanilla JavaScript on the CPU. This package is much smaller than the others because it doesn't need the TensorFlow binary, but it's also much slower.

Because this package doesn't rely on TensorFlow, it can be used in more devices that support Node.js. It's not limited to the Linux, Windows, and macOS platforms that support TensorFlow.

Production considerations

The Node.js bindings provide a backend for TensorFlow.js that implements operations synchronously. This means that, for example, when you call an operation like tf.matMul(a, b), it will block the main thread until the operation has completed.

For this reason, the bindings are well suited for scripts and offline tasks. If you want to use the Node.js bindings in a production application like a web server, you should set up a job queue or set up worker threads so your TensorFlow.js code doesn't block the main thread.

APIs

When you import the package as tf using any of the options above, all of the normal TensorFlow.js symbols appear on the imported module.

tf.browser

The APIs in the tf.browser.* namespace are not usable in Node.js because they depend on browser-specific APIs. For a list of the tf.browser APIs, see Browser.

tf.node

The two Node.js packages also provide a namespace, tf.node, which contains Node.js-specific APIs (for example, TensorBoard).

Here's an example of exporting summaries to TensorBoard in Node.js:

const model = tf.sequential();
model.add(tf.layers.dense({ units: 1, inputShape: [200] }));
model.compile({
  loss: 'meanSquaredError',
  optimizer: 'sgd',
  metrics: ['MAE']
});

// Generate some random fake data for demo purposes.
const xs = tf.randomUniform([10000, 200]);
const ys = tf.randomUniform([10000, 1]);
const valXs = tf.randomUniform([1000, 200]);
const valYs = tf.randomUniform([1000, 1]);

// Start model training process.
async function train() {
  await model.fit(xs, ys, {
    epochs: 100,
    validationData: [valXs, valYs],
    // Add the tensorBoard callback here.
    callbacks: tf.node.tensorBoard('/tmp/fit_logs_1')
  });
}
train();