title making new layers and models via subclassing fchollethttps twitter comfchollet date created 20190301 last modified 20230625 description complete guide to writing layer and model objects from scratch accelerator none introduction 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 the layer class the addweight method trainable and nontrainable weights the build method making sure your layers can be used with any backend the addloss method the training argument in call the mask argument in call making sure your layers can be serialized let s dive in setup the layer class the combination of state weights and some computation 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 denselyconnected layer it has two state variables the variables w and b you would use a layer by calling it on some tensor inputs much like a python function note that the weights w and b are automatically tracked by the layer upon being set as layer attributes layers can have nontrainable weights besides trainable weights you can add nontrainable 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 nontrainable weight it s part of layer weights but it gets categorized as a nontrainable weight it s not included in the trainable weights best practice deferring weight creation until the shape of the inputs is known our linear layer above took an inputdim 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 buildself inputsshape 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 at instantiation we don t know on what inputs this is going to get called the layer s weights are created dynamically the first time the layer is called implementing build separately as shown above nicely separates creating weights only once from using weights in every call layers are recursively composable if you assign a layer instance as an attribute of another layer the outer layer will start tracking the weights the inner layer we recommend creating such sublayers in the init method and leave it to the first call to trigger building their weights backendagnostic layers and backendspecific layers 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 the numpy api e g ops matmul ops sum ops reshape ops stack etc neural networksspecific apis such as ops softmax ops conv ops binarycrossentropy ops relu etc you can also use backendnative 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 jaxspecific layer using jax numpy python import jax class linearkeras layers layer def callself inputs return jax numpy matmulinputs self w self b this would be the equivalent tensorflowspecific layer python import tensorflow as tf class linearkeras layers layer def callself inputs return tf matmulinputs self w self b and this would be the equivalent pytorchspecific layer python import torch class linearkeras layers layer def callself inputs return torch matmulinputs self w self b because crossbackend compatibility is a tremendously useful property we strongly recommend that you seek to always make your layers backendagnostic by leveraging only keras apis the addloss method 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 addlossvalue a layer that creates an activity regularization loss these losses including those any inner layer can be retrieved via layer losses this property is reset at the start of every call to the toplevel layer so that layer losses always contains the loss values created during the last forward pass layer losses gets reset at the start of each call in addition the loss property also contains regularization losses created for the weights of any inner layer this is 1e3 sumlayer dense kernel 2 the kernelregularizer above 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 there is a loss passed in compile the regularization losses get added to it it s also possible not to pass any loss in compile since the model already has a loss to minimize via the addloss call during the forward pass you can optionally enable serialization on your layers if you need your custom layers to be serializable as part of a functional modelguidesfunctionalapi you can optionally implement a getconfig method now you can recreate the layer from its config 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 fromconfig class method this is the base implementation of fromconfig python def fromconfigcls config return clsconfig to learn more about serialization and saving see the complete guide to saving and serializing modelsguidesserializationandsaving privileged training argument in the call method 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 builtin training and evaluation loops e g fit to correctly use the layer in training and inference privileged mask argument in the call method 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 maskgenerating layers are the embedding layer configured with maskzerotrue and the masking layer the model class 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 the model class has the same api as layer with the following differences it exposes builtin training evaluation and prediction loops model fit model evaluate model predict it exposes the list of its inner layers via the model layers property it exposes saving and serialization apis save saveweights effectively the layer class corresponds to what we refer to in the literature as a layer as in convolution layer or recurrent layer or as a block as in resnet block or inception block meanwhile the model class corresponds to what is referred to in the literature as a model as in deep learning model or as a network as in deep neural network so if you re wondering should i use the layer class or the model class ask yourself will i need to call fit on it will i need to call save on it if so go with model if not either because your class is just a block in a bigger system or because you are writing training saving code yourself use layer for instance we could take our miniresnet example above and use it to build a model that we could train with fit and that we could save with saveweights python class resnetkeras model def initself numclasses1000 super init self block1 resnetblock self block2 resnetblock self globalpool layers globalaveragepooling2d self classifier densenumclasses def callself inputs x self block1inputs x self block2x x self globalpoolx return self classifierx resnet resnet dataset resnet fitdataset epochs10 resnet savefilepath keras putting it all together an endtoend example here s what you ve learned so far a layer encapsulate a state created in init or build and some computation defined in call layers can be recursively nested to create new bigger computation blocks layers are backendagnostic as long as they only use keras apis you can use backendnative apis such as jax numpy torch nn or tf nn but then your layer will only be usable with that specific backend layers can create and track losses typically regularization losses via addloss the outer container the thing you want to train is a model a model is just like a layer but with added training and serialization utilities let s put all of these things together into an endtoend example we re going to implement a variational autoencoder vae in a backendagnostic fashion so that it runs the same with tensorflow jax and pytorch we ll train it on mnist digits our vae will be a subclass of model built as a nested composition of layers that subclass layer it will feature a regularization loss kl divergence uses zmean zlogvar to sample z the vector encoding a digit def callself inputs zmean zlogvar inputs batch ops shapezmean0 dim ops shapezmean1 epsilon keras random normalshapebatch dim return zmean ops exp0 5 zlogvar epsilon class encoderlayers layer converts z the encoded digit vector back into a readable digit def init self originaldim intermediatedim64 namedecoder kwargs super initnamename kwargs self denseproj layers denseintermediatedim activationrelu self denseoutput layers denseoriginaldim activationsigmoid def callself inputs x self denseprojinputs return self denseoutputx class variationalautoencoderkeras model add kl divergence regularization loss let s train it on mnist using the fit api introduction 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 the layer class the add_weight method trainable and non trainable weights the build method making sure your layers can be used with any backend the add_loss method the training argument in call the mask argument in call making sure your layers can be serialized let s dive in setup the layer class the combination of state weights and some computation 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 layers can have non trainable weights 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 it s not included in the trainable weights best practice deferring weight creation until the shape of the inputs is known 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 at instantiation we don t know on what inputs this is going to get called the layer s weights are created dynamically the first time the layer is called implementing build separately as shown above nicely separates creating weights only once from using weights in every call layers are recursively composable if you assign a layer instance as an attribute of another layer the outer layer will start tracking the weights the inner layer we recommend creating such sublayers in the __init__ method and leave it to the first __call__ to trigger building their weights the first call to the mlp will create the weights backend agnostic layers and backend specific layers 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 the numpy api e g ops matmul ops sum ops reshape ops stack etc neural networks specific apis such as ops softmax ops conv ops binary_crossentropy ops relu etc 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 python import jax class linear keras layers layer def call self inputs return jax numpy matmul inputs self w self b this would be the equivalent tensorflow specific layer python import tensorflow as tf class linear keras layers layer def call self inputs return tf matmul inputs self w self b and this would be the equivalent pytorch specific layer python import torch class linear keras layers layer def call self inputs return torch matmul inputs self w self b 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 the add_loss method 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 a layer that creates an activity regularization loss these losses including those 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 no losses yet since the layer has never been called we created one loss value layer losses gets reset at the start of each __call__ this is the loss created during the call above in addition the loss property also contains regularization losses created for the weights of any inner layer this is 1e 3 sum layer dense kernel 2 the kernel_regularizer above 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 there is a loss passed in compile the regularization losses get added to it it s also possible not to pass any loss in compile since the model already has a loss to minimize via the add_loss call during the forward pass you can optionally enable serialization on your layers if you need your custom layers to be serializable as part of a functional model guides functional_api you can optionally implement a get_config method now you can recreate the layer from its config 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 python def from_config cls config return cls config to learn more about serialization and saving see the complete guide to saving and serializing models guides serialization_and_saving privileged training argument in the call method 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 privileged mask argument in the call method 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 the model class 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 the model class has the same api as layer with the following differences it exposes built in training evaluation and prediction loops model fit model evaluate model predict it exposes the list of its inner layers via the model layers property it exposes saving and serialization apis save save_weights effectively the layer class corresponds to what we refer to in the literature as a layer as in convolution layer or recurrent layer or as a block as in resnet block or inception block meanwhile the model class corresponds to what is referred to in the literature as a model as in deep learning model or as a network as in deep neural network so if you re wondering should i use the layer class or the model class ask yourself will i need to call fit on it will i need to call save on it if so go with model if not either because your class is just a block in a bigger system or because you are writing training saving code yourself use layer for instance we could take our mini resnet example above and use it to build a model that we could train with fit and that we could save with save_weights python class resnet keras model def __init__ self num_classes 1000 super __init__ self block_1 resnetblock self block_2 resnetblock self global_pool layers globalaveragepooling2d self classifier dense num_classes def call self inputs x self block_1 inputs x self block_2 x x self global_pool x return self classifier x resnet resnet dataset resnet fit dataset epochs 10 resnet save filepath keras putting it all together an end to end example here s what you ve learned so far a layer encapsulate a state created in __init__ or build and some computation defined in call layers can be recursively nested to create new bigger computation blocks layers are backend agnostic as long as they only use keras apis you can use backend native apis such as jax numpy torch nn or tf nn but then your layer will only be usable with that specific backend layers can create and track losses typically regularization losses via add_loss the outer container the thing you want to train is a model a model is just like a layer but with added training and serialization utilities let s put all of these things together into an end to end example we re going to implement a variational autoencoder vae in a backend agnostic fashion so that it runs the same with tensorflow jax and pytorch we ll train it on mnist digits our vae will be a subclass of model built as a nested composition of layers that subclass layer it will feature a regularization loss kl divergence uses z_mean z_log_var to sample z the vector encoding a digit maps mnist digits to a triplet z_mean z_log_var z converts z the encoded digit vector back into a readable digit combines the encoder and decoder into an end to end model for training add kl divergence regularization loss let s train it on mnist using the fit api
import numpy as np
import keras
from keras import ops
from keras import layers




class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super().__init__()
        self.w = self.add_weight(
            shape=(input_dim, units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(units,), initializer="zeros", trainable=True
        )

    def call(self, inputs):
        return ops.matmul(inputs, self.w) + self.b




x = ops.ones((2, 2))
linear_layer = Linear(4, 2)
y = linear_layer(x)
print(y)



assert linear_layer.weights == [linear_layer.w, linear_layer.b]




class ComputeSum(keras.layers.Layer):
    def __init__(self, input_dim):
        super().__init__()
        self.total = self.add_weight(
            initializer="zeros", shape=(input_dim,), trainable=False
        )

    def call(self, inputs):
        self.total.assign_add(ops.sum(inputs, axis=0))
        return self.total


x = ops.ones((2, 2))
my_sum = ComputeSum(2)
y = my_sum(x)
print(y.numpy())
y = my_sum(x)
print(y.numpy())



print("weights:", len(my_sum.weights))
print("non-trainable weights:", len(my_sum.non_trainable_weights))


print("trainable_weights:", my_sum.trainable_weights)




class Linear(keras.layers.Layer):
    def __init__(self, units=32, input_dim=32):
        super().__init__()
        self.w = self.add_weight(
            shape=(input_dim, units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(units,), initializer="zeros", trainable=True
        )

    def call(self, inputs):
        return ops.matmul(inputs, self.w) + self.b





class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super().__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,), initializer="random_normal", trainable=True
        )

    def call(self, inputs):
        return ops.matmul(inputs, self.w) + self.b





linear_layer = Linear(32)


y = linear_layer(x)






class MLPBlock(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.linear_1 = Linear(32)
        self.linear_2 = Linear(32)
        self.linear_3 = Linear(1)

    def call(self, inputs):
        x = self.linear_1(inputs)
        x = keras.activations.relu(x)
        x = self.linear_2(x)
        x = keras.activations.relu(x)
        return self.linear_3(x)


mlp = MLPBlock()
y = mlp(
    ops.ones(shape=(3, 64))
)  
print("weights:", len(mlp.weights))
print("trainable weights:", len(mlp.trainable_weights))







class ActivityRegularizationLayer(keras.layers.Layer):
    def __init__(self, rate=1e-2):
        super().__init__()
        self.rate = rate

    def call(self, inputs):
        self.add_loss(self.rate * ops.mean(inputs))
        return inputs





class OuterLayer(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.activity_reg = ActivityRegularizationLayer(1e-2)

    def call(self, inputs):
        return self.activity_reg(inputs)


layer = OuterLayer()
assert (
    len(layer.losses) == 0
)  

_ = layer(ops.zeros((1, 1)))
assert len(layer.losses) == 1  


_ = layer(ops.zeros((1, 1)))
assert len(layer.losses) == 1  




class OuterLayerWithKernelRegularizer(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.dense = keras.layers.Dense(
            32, kernel_regularizer=keras.regularizers.l2(1e-3)
        )

    def call(self, inputs):
        return self.dense(inputs)


layer = OuterLayerWithKernelRegularizer()
_ = layer(ops.zeros((1, 1)))



print(layer.losses)



inputs = keras.Input(shape=(3,))
outputs = ActivityRegularizationLayer()(inputs)
model = keras.Model(inputs, outputs)



model.compile(optimizer="adam", loss="mse")
model.fit(np.random.random((2, 3)), np.random.random((2, 3)))




model.compile(optimizer="adam")
model.fit(np.random.random((2, 3)), np.random.random((2, 3)))




class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super().__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,), initializer="random_normal", trainable=True
        )

    def call(self, inputs):
        return ops.matmul(inputs, self.w) + self.b

    def get_config(self):
        return {"units": self.units}



layer = Linear(64)
config = layer.get_config()
print(config)
new_layer = Linear.from_config(config)




class Linear(keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,), initializer="random_normal", trainable=True
        )

    def call(self, inputs):
        return ops.matmul(inputs, self.w) + self.b

    def get_config(self):
        config = super().get_config()
        config.update({"units": self.units})
        return config


layer = Linear(64)
config = layer.get_config()
print(config)
new_layer = Linear.from_config(config)






class CustomDropout(keras.layers.Layer):
    def __init__(self, rate, **kwargs):
        super().__init__(**kwargs)
        self.rate = rate

    def call(self, inputs, training=None):
        if training:
            return keras.random.dropout(inputs, rate=self.rate)
        return inputs











class Sampling(layers.Layer):
    

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = ops.shape(z_mean)[0]
        dim = ops.shape(z_mean)[1]
        epsilon = keras.random.normal(shape=(batch, dim))
        return z_mean + ops.exp(0.5 * z_log_var) * epsilon


class Encoder(layers.Layer):
    

    def __init__(
        self, latent_dim=32, intermediate_dim=64, name="encoder", **kwargs
    ):
        super().__init__(name=name, **kwargs)
        self.dense_proj = layers.Dense(intermediate_dim, activation="relu")
        self.dense_mean = layers.Dense(latent_dim)
        self.dense_log_var = layers.Dense(latent_dim)
        self.sampling = Sampling()

    def call(self, inputs):
        x = self.dense_proj(inputs)
        z_mean = self.dense_mean(x)
        z_log_var = self.dense_log_var(x)
        z = self.sampling((z_mean, z_log_var))
        return z_mean, z_log_var, z


class Decoder(layers.Layer):
    

    def __init__(
        self, original_dim, intermediate_dim=64, name="decoder", **kwargs
    ):
        super().__init__(name=name, **kwargs)
        self.dense_proj = layers.Dense(intermediate_dim, activation="relu")
        self.dense_output = layers.Dense(original_dim, activation="sigmoid")

    def call(self, inputs):
        x = self.dense_proj(inputs)
        return self.dense_output(x)


class VariationalAutoEncoder(keras.Model):
    

    def __init__(
        self,
        original_dim,
        intermediate_dim=64,
        latent_dim=32,
        name="autoencoder",
        **kwargs
    ):
        super().__init__(name=name, **kwargs)
        self.original_dim = original_dim
        self.encoder = Encoder(
            latent_dim=latent_dim, intermediate_dim=intermediate_dim
        )
        self.decoder = Decoder(original_dim, intermediate_dim=intermediate_dim)

    def call(self, inputs):
        z_mean, z_log_var, z = self.encoder(inputs)
        reconstructed = self.decoder(z)
        
        kl_loss = -0.5 * ops.mean(
            z_log_var - ops.square(z_mean) - ops.exp(z_log_var) + 1
        )
        self.add_loss(kl_loss)
        return reconstructed




(x_train, _), _ = keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype("float32") / 255

original_dim = 784
vae = VariationalAutoEncoder(784, 64, 32)

optimizer = keras.optimizers.Adam(learning_rate=1e-3)
vae.compile(optimizer, loss=keras.losses.MeanSquaredError())

vae.fit(x_train, x_train, epochs=2, batch_size=64)