This tutorial is an adaptation of Francois Chollet's brilliant introduction to TensorFlow 2.0. It modifies the code to add the Weights & Biases callbacks to help the reader save and track the progress of their models.

You find the accompanying W&B dashboard here.

TensorFlow 2.0 + Keras Overview for Deep Learning Researchers

@fchollet, October 2019

This document serves as an introduction, crash course, and quick API reference for TensorFlow 2.0.

TensorFlow and Keras were both released over four years ago (March 2015 for Keras and November 2015 for TensorFlow). That's a long time in deep learning years!

In the old days, TensorFlow 1.x + Keras had a number of known issues:

• Using TensorFlow meant manipulating static computation graphs, which would feel awkward and difficult to programmers used to imperative styles of coding.
• While the TensorFlow API was very powerful and flexible, it lacked polish and was often confusing or difficult to use.
• While Keras was very productive and easy to use, it would often lack flexibility for research use cases.

TensorFlow 2.0 is an extensive redesign of TensorFlow and Keras that takes into account over four years of user feedback and technical progress. It fixes the issues above in a big way.

It's a machine learning platform from the future.

TensorFlow 2.0 is built on the following key ideas:

• Let users run their computation eagerly, like they would in Numpy. This makes TensorFlow 2.0 programming intuitive and Pythonic.
• Preserve the considerable advantages of compiled graphs (for performance, distribution, and deployment). This makes TensorFlow fast, scalable, and production-ready.
• Leverage Keras as its high-level deep learning API, making TensorFlow approachable and highly productive.
• Extend Keras into a spectrum of workflows ranging from the very high-level (easier to use, less flexible) to the very low-level (requires more expertise, but provides great flexibility).

Part 1: TensorFlow basics

Tensors

This is a constant tensor:

You can get its value as a Numpy array by calling .numpy():

Much like a Numpy array, it features the attributes dtype and shape:

A common way to create constant tensors is via tf.ones and tf.zeros (just like np.ones and np.zeros):

Random constant tensors

This is all pretty normal:

And here's an integer tensor with values drawn from a random uniform distribution:

Variables

Variables are special tensors used to store mutable state (like the weights of a neural network). You create a Variable using some initial value.

You update the value of a Variable by using the methods .assign(value), or .assign_add(increment) or .assign_sub(decrement):

Doing math in TensorFlow

You can use TensorFlow exactly like you would use Numpy. The main difference is that your TensorFlow code can run on GPU and TPU.

Computing gradients with GradientTape

Oh, and there's another big difference with Numpy: you can automatically retrieve the gradient of any differentiable expression.

Just open a GradientTape, start "watching" a tensor via tape.watch(), and compose a differentiable expression using this tensor as input:

By default, variables are watched automatically, so you don't need to manually watch them:

Note that you can compute higher-order derivatives by nesting tapes:

An end-to-end example: linear regression

So far you've learned that TensorFlow is a Numpy-like library that is GPU or TPU accelerated, with automatic differentiation. Time for an end-to-end example: let's implement a linear regression, the FizzBuzz of Machine Learning.

For the sake of demonstration, we won't use any of the higher-level Keras components like Layer or MeanSquaredError. Just basic ops.

Let's generate some artificial data to demonstrate our model:

Now let's train our linear regression by iterating over batch-by-batch over the data and repeatedly calling train_on_batch:

Here's how our model performs:

Making it fast with tf.function

But how fast is our current code running?

Let's compile the training function into a static graph. Literally all we need to do is add the tf.function decorator on it:

Let's try this again:

40% reduction, neat. In this case we used a trivially simple model; in general the bigger the model the greater the speedup you can get by leveraging static graphs.

Remember: eager execution is great for debugging and printing results line-by-line, but when it's time to scale, static graphs are a researcher's best friends.

Part 2: The Keras API

Keras is a Python API for deep learning. It has something for everyone:

• If you're an engineer, Keras provides you with reusable blocks such as layers, metrics, training loops, to support common use cases. It provides a high-level user experience that's accessible and productive.

• If you're a researcher, you may prefer not to use these built-in blocks such as layers and training loops, and instead create your own. Of course, Keras allows you to do this. In this case, Keras provides you with templates for the blocks you write, it provides you with structure, with an API standard for things like Layers and Metrics. This structure makes your code easy to share with others and easy to integrate in production workflows.

• The same is true for library developers: TensorFlow is a large ecosystem. It has many different libraries. In order for different libraries to be able to talk to each other and share components, they need to follow an API standard. That's what Keras provides.

Crucially, Keras brings high-level UX and low-level flexibility together fluently: you no longer have on one hand, a high-level API that's easy to use but inflexible, and on the other hand a low-level API that's flexible but only approachable by experts. Instead, you have a spectrum of workflows, from the very high-level to the very low-level. Workflows that are all compatible because they're built on top of the same concepts and objects.

The base Layer class

The first class you need to know is Layer. Pretty much everything in Keras derives from it.

A Layer encapsulates a state (weights) and some computation (defined in the call method).

A layer instance works like a function. Let's call it on some data:

The Layer class takes care of tracking the weights assigned to it as attributes:

Note that's also a shortcut method for creating weights: add_weight. Instead of doing

	w_init = tf.random_normal_initializer()
self.w = tf.Variable(initial_value=w_init(shape=shape, dtype='float32'))

You would typically do:

	self.w = self.add_weight(shape=shape, initializer='random_normal')

It’s good practice to create weights in a separate build method, called lazily with the shape of the first inputs seen by your layer. Here, this pattern prevents us from having to specify input_dim in the constructor:

Trainable and non-trainable weights

Weights created by layers can be either trainable or non-trainable. They're exposed in trainable_weights and non_trainable_weights. Here's a layer with a non-trainable weight:

Recursively composing layers

Layers can be recursively nested to create bigger computation blocks. Each layer will track the weights of its sublayers (both trainable and non-trainable.

Built-in layers

Keras provides you with a wide range of built-in layers, so that you don't have to implement your own layers all the time.

• Convolution layers
• Transposed convolutions
• Separateable convolutions
• Average and max pooling
• Global average and max pooling
• LSTM, GRU (with built-in cuDNN acceleration)
• BatchNormalization
• Dropout
• Attention
• ConvLSTM2D
• etc.

Keras follows the principles of exposing good default configurations, so that layers will work fine out of the box for most use cases if you leave keyword arguments to their default value. For instance, the LSTM layer uses an orthogonal recurrent matrix initializer by default, and initializes the forget gate bias to one by default.

The training argument in call

Some layers, in particular the BatchNormalization layer and the Dropout layer, have different behaviors during training and inference. For such layers, it is standard practice to expose a training (boolean) argument in the call method.

By exposing this argument in call, you enable the built-in training and evaluation loops (e.g. fit) to correctly use the layer in training and inference.

A more Functional way of defining models

To build deep learning models, you don't have to use object-oriented programming all the time. Layers can also be composed functionally, like this (we call it the "Functional API"):

The Functional API tends to be more concise than subclassing, and provides a few other advantages (generally the same advantages that functional, typed languages provide over untyped OO development). However, it can only be used to define DAGs of layers -- recursive networks should be defined as Layer subclasses instead.

Key differences between models defined via subclassing and Functional models are explained in this blog post.

Learn more about the Functional API here.

In your research workflows, you may often find yourself mix-and-matching OO models and Functional models.

For models that are simple stacks of layers with a single input and a single output, you can also use the Sequential class which turns a list of layers into a Model:

Loss classes

Keras features a wide range of built-in loss classes, like BinaryCrossentropy, CategoricalCrossentropy, KLDivergence, etc. They work like this:

Note that loss classes are stateless: the output of __call__ is only a function of the input.

Metric classes

Keras also features a wide range of built-in metric classes, such as BinaryAccuracy, AUC, FalsePositives, etc.

Unlike losses, metrics are stateful. You update their state using the update_state method, and you query the scalar metric result using result:

The internal state can be cleared with metric.reset_states.

You can easily roll out your own metrics by subclassing the Metric class:

• Create the state variables in __init__
• Update the variables given y_true and y_pred in update_state
• Return the metric result in result
• Clear the state in reset_states

Here's a quick implementation of a BinaryTruePositives metric as a demonstration:

Optimizer classes & a quick end-to-end training loop

You don't normally have to define by hand how to update your variables during gradient descent, like we did in our initial linear regression example. You would usually use one of the built-in Keras optimizer, like SGD, RMSprop, or Adam.

Here's a simple MNSIT example that brings together loss classes, metric classes, and optimizers.

We can reuse our SparseCategoricalAccuracy metric instance to implement a testing loop:

The add_loss method

Sometimes you need to compute loss values on the fly during a foward pass (especially regularization losses). Keras allows you to compute loss values at any time, and to recursively keep track of them via the add_loss method.

Here's an example of a layer that adds a sparsity regularization loss based on the L2 norm of the inputs:

Loss values added via add_loss can be retrieved in the .losses list property of any Layer or Model:

These losses are cleared by the top-level layer at the start of each forward pass -- they don't accumulate. So layer.losses always contain only the losses created during the last forward pass. You would typically use these losses by summing them before computing your gradients when writing a training loop.

A detailed end-to-end example: a Variational AutoEncoder (VAE)

If you want to take a break from the basics and look at a slightly more advanced example, check out this Variational AutoEncoder implementation that demonstrates everything you've learned so far:

• Subclassing Layer
• Recursive layer composition
• Loss classes and metric classes
• add_loss
• GradientTape

Using built-in training loops

It would be a bit silly if you had to write your own low-level training loops every time for simple use cases. Keras provides you with a built-in training loop on the Model class. If you want to use it, either subclass from Model or create a Functional or Sequential model.

To demonstrate it, let's reuse the MNIST setup from above:

First, call compile to configure the optimizer, loss, and metrics to monitor.

Then we call fit on our model to pass it the data:

Done! Now let's test it:

Note that you can also monitor your loss and metrics on some validation data during fit.

Also, you can call fit directly on Numpy arrays, so no need for the dataset conversion:

Callbacks

One of the neat features of fit (besides built-in support for sample weighting and class weighting) is that you can easily customize what happens during training and evaluation by using callbacks.

A callback is an object that is called at different points during training (e.g. at the end of every batch or at the end of every epoch) and does stuff.

There's a bunch of built-in callback available, like ModelCheckpoint to save your models after each epoch during training, or EarlyStopping, which interrupts training when your validation metrics start stalling.

And you can easily write your own callbacks.

Parting words

I hope this guide has given you a good overview of what's possible with TensorFlow 2.0 and Keras!

Remember that TensorFlow and Keras don't represent a single workflow. It's a spectrum of workflows, each with its own trade-off between usability and flexbility. For instance, you've noticed that it's much easier to use fit than to write a custom training loop, but fit doesn't give you the same level of granular control for research use cases.

So use the right tool for the job!

A core principle of Keras is "gradual disclosure of complexity": it's easy to get started, and you can gradually dive into workflows where you write more and more logic from scratch, providing you with complete control.

This applies to both model definition, and model training.

What to learn next

Next, there are many more topics you may be interested in:

• Saving and serialization
• Distributed training on multiple GPUS
• Exporting models to TFLite for deployment on Android or embedded systems
• Exporting models to TensorFlow.js for deployment in the browser