TensorFlow.js is an open source WebGL-accelerated JavaScript library for machine intelligence. It brings highly performant machine learning building blocks to your fingertips, allowing you to train neural networks in a browser or run pre-trained models in inference mode. See Getting Started for a guide on installing/configuring TensorFlow.js.
TensorFlow.js provides low-level building blocks for machine learning as well as a high-level, Keras-inspired API for constructing neural networks. Let's take a look at some of the core components of the library.
The central unit of data in TensorFlow.js is the tensor: a set of numerical values shaped into an array of one or more dimensions. A Tensor instance has a shape attribute that defines the array shape (i.e., how many values are in each dimension of the array).
The primary Tensor constructor is the tf.tensor function:
// 2x3 Tensor
→ const shape = [2, 3]; // 2 rows, 3 columns
→ const a = tf.tensor([1.0, 2.0, 3.0, 10.0, 20.0, 30.0], shape);
a.print(); // print Tensor values
// Output: [[1 , 2 , 3 ],
// [10, 20, 30]]
// The shape can also be inferred:
const b = tf.tensor([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]);
b.print();
// Output: [[1 , 2 , 3 ],
// [10, 20, 30]]
However, for constructing low-rank tensors, we recommend using the following
functions to enhance code readability: tf.scalar, tf.tensor1d, tf.tensor2d, tf.tensor3d and tf.tensor4d.
The following example creates an identical tensor to the one above using tf.tensor2d:
→ const c = tf.tensor2d([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]);
c.print();
// Output: [[1 , 2 , 3 ],
// [10, 20, 30]]
TensorFlow.js also provides convenience functions for creating tensors with all values
set to 0 (tf.zeros) or all values set to 1 (tf.ones):
// 3x5 Tensor with all values set to 0
→ const zeros = tf.zeros([3, 5]);
// Output: [[0, 0, 0, 0, 0],
// [0, 0, 0, 0, 0],
// [0, 0, 0, 0, 0]]
In TensorFlow.js, tensors are immutable; once created, you cannot change their values. Instead you perform operations on them that generate new tensors.
Variables are initialized with a tensor of values. Unlike Tensors, however, their values are mutable. You can assign a new tensor to an existing variable using the assign method:
const initialValues = tf.zeros([5]);
→ const biases = tf.variable(initialValues); // initialize biases
biases.print(); // output: [0, 0, 0, 0, 0]
const updatedValues = tf.tensor1d([0, 1, 0, 1, 0]);
biases.assign(updatedValues); // update values of biases
biases.print(); // output: [0, 1, 0, 1, 0]
Variables are primarily used to store and then update values during model training.
While tensors allow you to store data, operations (ops) allow you to manipulate that data. TensorFlow.js provides a wide variety of ops suitable for linear algebra and machine learning that can be performed on tensors. Because tensors are immutable, these ops do not change their values; instead, ops return new tensors.
Available ops include unary ops such as square:
const d = tf.tensor2d([[1.0, 2.0], [3.0, 4.0]]);
→ const d_squared = d.square();
d_squared.print();
// Output: [[1, 4 ],
// [9, 16]]
And binary ops such as add, sub, and mul:
const e = tf.tensor2d([[1.0, 2.0], [3.0, 4.0]]);
const f = tf.tensor2d([[5.0, 6.0], [7.0, 8.0]]);
const e_plus_f = e.add(f);
e_plus_f.print();
// Output: [[6 , 8 ],
// [10, 12]]
TensorFlow.js has a chainable API; you can call ops on the result of ops:
const sq_sum = e.add(f).square();
sq_sum.print();
// Output: [[36 , 64 ],
// [100, 144]]
// All operations are also exposed as functions in the main namespace,
// so you could also do the following:
→ const sq_sum = tf.square(tf.add(e, f));
Conceptually, a model is a function that given some input will produce some desired output.
In TensorFlow.js there are two ways to create models. You can use ops directly to represent the work the model does. For example:
// → Define function
function predict(input) {
// y = a * x ^ 2 + b * x + c
// More on tf.tidy in the next section
return tf.tidy(() => {
const x = tf.scalar(input);
const ax2 = a.mul(x.square());
const bx = b.mul(x);
const y = ax2.add(bx).add(c);
return y;
});
}
// Define constants: y = 2x^2 + 4x + 8
const a = tf.scalar(2);
const b = tf.scalar(4);
const c = tf.scalar(8);
// Predict output for input of 2
const result = predict(2);
result.print() // Output: 24
→ You can also use the high-level API tf.model to construct a model out of layers, which are a popular abstraction in deep learning.
The TensorFlow layers module provides a high-level API that makes
it easy to construct a neural network. It provides methods that facilitate the
creation of dense (fully connected) layers and convolutional layers, adding
activation functions, and applying dropout regularization.
The following code constructs a tf.sequential model:
→ const model = tf.sequential();
model.add(
→ tf.layers.simpleRNN({
units: 20,
recurrentInitializer: 'GlorotNormal',
inputShape: [80, 4]
})
);
→ const optimizer = tf.train.sgd(LEARNING_RATE);
→ model.compile({optimizer, loss: 'categoricalCrossentropy'});
→ model.fit({x: data, y: labels});
There are many different types of layers available in TensorFlow.js.
→ A few examples include tf.layers.simpleRNN, tf.layers.gru, and tf.layers.lstm.
Because TensorFlow.js uses the GPU to accelerate math operations, it's necessary to manage GPU memory when working with tensors and variables.
TensorFlow.js provide two functions to help with this: dispose and tf.tidy.
→ You can call dispose on a tensor or variable to purge it and free up its GPU memory:
const x = tf.tensor2d([[0.0, 2.0], [4.0, 6.0]]);
const x_squared = x.square();
→ x.dispose();
x_squared.dispose();
Using dispose can be cumbersome when doing a lot of tensor operations. TensorFlow.js provides another function, tf.tidy, that plays a similar role to regular scopes in JavaScript, but for GPU-backed tensors.
→ tf.tidy executes a function and purges any intermediate tensors created, freeing up their GPU memory. It does not purge the return value of the inner function.
// tf.tidy takes a function to tidy up after
→ const average = tf.tidy(() => {
// tf.tidy will clean up all the GPU memory used by tensors inside
// this function, other than the tensor that is returned.
//
// Even in a short sequence of operations like the one below, a number
// of intermediate tensors get created. So it is a good practice to
// put your math ops in a tidy!
const y = tf.tensor1d([1.0, 2.0, 3.0, 4.0]);
const z = tf.ones([4]);
return y.sub(z).square().mean();
});
average.print() // Output: 3.5
Using tf.tidy will help prevent memory leaks in your application. It can also be used to more carefully control when memory is reclaimed.
The function passed to tf.tidy should be synchronous and also not return a Promise. We suggest keeping code that updates the UI or makes remote requests outside of tf.tidy.
→ tf.tidy will not clean up variables. Variables typically last through the entire lifecycle of a machine learning model, so TensorFlow.js doesn't clean them up even if they are created in a tidy; however, you can call dispose on them manually.
See the TensorFlow.js API reference for comprehensive documentation of the library.
For a more in-depth look at machine learning fundamentals, see the following resources: