Making new layers and models via subclassing

Author: fchollet Date created: 2019/03/01 Last modified: 2023/06/25 Description: Complete guide to writing Layer and Model objects from scratch.
View in Colab • GitHub source
This guide will cover everything you need to know to build your own subclassed layers and models. In particular, you'll learn about the following features:
One of the central abstractions in Keras is the Layer class. A layer encapsulates both a state (the layer's "weights") and a transformation from inputs to outputs (a "call", the layer's forward pass).
Here's a densely-connected layer. It has two state variables: the variables w and b.
You would use a layer by calling it on some tensor input(s), much like a Python function.
Note that the weights w and b are automatically tracked by the layer upon being set as layer attributes:
Besides trainable weights, you can add non-trainable weights to a layer as well. Such weights are meant not to be taken into account during backpropagation, when you are training the layer.
Here's how to add and use a non-trainable weight:
It's part of layer.weights, but it gets categorized as a non-trainable weight:
Our Linear layer above took an input_dim argument that was used to compute the shape of the weights w and b in __init__():
In many cases, you may not know in advance the size of your inputs, and you would like to lazily create weights when that value becomes known, some time after instantiating the layer.
In the Keras API, we recommend creating layer weights in the build(self, inputs_shape) method of your layer. Like this:
The __call__() method of your layer will automatically run build the first time it is called. You now have a layer that's lazy and thus easier to use:
Implementing build() separately as shown above nicely separates creating weights only once from using weights in every call.
If you assign a Layer instance as an attribute of another Layer, the outer layer will start tracking the weights created by the inner layer.
We recommend creating such sublayers in the __init__() method and leave it to the first __call__() to trigger building their weights.
As long as a layer only uses APIs from the keras.ops namespace (or other Keras namespaces such as keras.activations, keras.random, or keras.layers), then it can be used with any backend – TensorFlow, JAX, or PyTorch.
All layers you've seen so far in this guide work with all Keras backends.
The keras.ops namespace gives you access to:
You can also use backend-native APIs in your layers (such as tf.nn functions), but if you do this, then your layer will only be usable with the backend in question. For instance, you could write the following JAX-specific layer using jax.numpy:
This would be the equivalent TensorFlow-specific layer:
And this would be the equivalent PyTorch-specific layer:
Because cross-backend compatibility is a tremendously useful property, we strongly recommend that you seek to always make your layers backend-agnostic by leveraging only Keras APIs.
When writing the call() method of a layer, you can create loss tensors that you will want to use later, when writing your training loop. This is doable by calling self.add_loss(value):
These losses (including those created by any inner layer) can be retrieved via layer.losses. This property is reset at the start of every __call__() to the top-level layer, so that layer.losses always contains the loss values created during the last forward pass.
In addition, the loss property also contains regularization losses created for the weights of any inner layer:
These losses are meant to be taken into account when writing custom training loops.
They also work seamlessly with fit() (they get automatically summed and added to the main loss, if any):
If you need your custom layers to be serializable as part of a Functional model, you can optionally implement a get_config() method:
Note that the __init__() method of the base Layer class takes some keyword arguments, in particular a name and a dtype. It's good practice to pass these arguments to the parent class in __init__() and to include them in the layer config:
If you need more flexibility when deserializing the layer from its config, you can also override the from_config() class method. This is the base implementation of from_config():
To learn more about serialization and saving, see the complete guide to saving and serializing models.
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.
The other privileged argument supported by call() is the mask argument.
You will find it in all Keras RNN layers. A mask is a boolean tensor (one boolean value per timestep in the input) used to skip certain input timesteps when processing timeseries data.
Keras will automatically pass the correct mask argument to __call__() for layers that support it, when a mask is generated by a prior layer. Mask-generating layers are the Embedding layer configured with mask_zero=True, and the Masking layer.
In general, you will use the Layer class to define inner computation blocks, and will use the Model class to define the outer model – the object you will train.
For instance, in a ResNet50 model, you would have several ResNet blocks subclassing Layer, and a single Model encompassing the entire ResNet50 network.