Building Custom Layers in Keras

Creating a custom layer in Keras

Keras provides us with a pre-defined layer class to create layers. But this will limit our scope to already defined logic and functionalities. Thus creating our own custom layer will provide us with a wide range of options and full control over their functionalities. 

Now let’s start implementing our custom layer in Keras.

Importing libraries

First, let’s start by importing the necessary libraries, which will be required to build the layer.

import tensorflow as tf
from tensorflow import keras
from keras.layers import Layer

 

Here, the first line imports TensorFlow from which the Keras module is imported. This will allow us to access various functionalities of the TensorFlow library. And lastly, the Layer class is imported, which is the base class for building custom layers in Keras.

Defining the Custom Layer Class

Let us now create a new custom layer class called My_CustomLayer. This class will inherit from the ‘Layer’ class. By subclassing it, we will create custom layers with specific functionalities.

class My_CustomLayer(layers.Layer):
…
… 

Initialising Custom Layer Class

Let us initialise our new class as shown below:

def __init__(self, output_dim, units=32, **kwargs):
	super().__init__(**kwargs)
	self.units = units
	self.output_dim = output_dim
	self.kernel = None

 

Here,

• def__init__(): It is the constructor method that takes two required arguments, and the **kwargs argument allows to accept additional arguments.
   
• super().__init__(): It calls the constructor of the base class for proper initialisation of the layer class.
   
• The rest lines assign the attribute values of the layer.

Implementing Build Method

The build method is used to create the weights of the layers. Its only purpose is to build the layers properly once called. Below is our custom-build method:

def build(self, input_shape):
	self.kernel = self.add_weight(
		name='kernel',
		shape=(input_shape[1], self.output_dim),
		initializer=tf.random_normal_initializer(),
		trainable=True
	)
	super().build(input_shape)

 

Here,

• self.kernal = self.add_weight (...): This line creates a weight variable called kernel with the defined shape of weights. Further, the weights are initialised with random normal distribution.
   
• super().build(input_shape): This line calls the ‘build’ method of the base ‘Layer’ class.

Implementing call Method

Below is the implementation of the custom call method:

def call(self, inputs):
	return tf.matmul(inputs, self.kernel)

 

Here, 

• def call(...): This line creates a call method with one argument, ‘inputs’, which is the input data for our layer.
   
• return tf.matmul(...): It performs and returns the matrix multiplication between input and self.kernel.

Implementing compute_output_shape Method

Let us now create the method to calculate the output shape:

def compute_output_shape(self, input_shape):
	return (input_shape[0], self.output_dim)

 

This method will compute the output shape of the layer based on the input shape and returns a tuple as output.

Implementing get_config Method

Below is the implementation of the get_config() method:

def get_config(self):
	return {'output_dim': self.output_dim, 'units': self.units}

 

This method returns the configuration of our custom layer as a dictionary.

The final code of our custom layer is:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

class My_CustomLayer(layers.Layer):
	def __init__(self, output_dim, units=32, **kwargs):
		super().__init__(**kwargs)
		self.units = units
		self.output_dim = output_dim
		self.kernel = None

	def build(self, input_shape):
		self.kernel = self.add_weight(
			name='kernel',
			shape=(input_shape[1], self.output_dim),
			initializer=tf.random_normal_initializer(),
			trainable=True
		)
		super().build(input_shape)

	def call(self, inputs):
		return tf.matmul(inputs, self.kernel)

	def compute_output_shape(self, input_shape):
		return (input_shape[0], self.output_dim)

	def get_config(self):
		return {'output_dim': self.output_dim, 'units': self.units}

 

Using Our Custom Layer

Now, let us create a simple model by using our custom layer.

#Importing libraries
from keras.models import Sequential 
from keras.layers import Dense 

#Creating a simple sequential model
model = Sequential() 
model.add(My_CustomLayer(16, input_shape = (16,))) 
model.add(Dense(8, activation = 'softmax')) 
model.add(Dense(2, input_dim=1, activation='relu'))
model.add(My_CustomLayer(32, input_shape = (16,))) 
model.add(Dense(1, activation='sigmoid'))
model.add(Dense(2, activation='softmax'))

#Calling the summary method
model.summary()

 

Here, we created a simple sequential model and added multiple layers to it. We also added two custom layers with their respective units and input shapes and called the summary method to give a total overview of the created model.

The output after executing the above code will be:

Thus, we have successfully built a deep learning model with our custom layer in Keras.

Frequently Asked Questions

What is Keras?

Keras is an open-source deep-learning library in Python. It provides high-level neural network API to build deep learning models.

What is a custom layer in Keras?

In Keras, a custom layer is a user-defined layer that inherits properties from the base ‘Layer’ class.

Are custom layers compatible with all model types?

Yes, the custom layers in Keras are compatible with functional API and subclassed models of any type.

How can I view detail about my model in Keras?

The summary() method in Keras is used to get a detailed overview of a model or data frame.

Conclusion

This article discusses the topic of building custom layers in Keras. We discussed the Layer class in Keras, using which we created a custom layer. We hope this blog has helped you enhance your knowledge of building custom layers in Keras. If you want to learn more, then check out our articles.

• Application of Keras
• Transfer Learning with Keras
• Visualizing Vision Models in Keras
   

Refer to our Guided Path to upskill yourself in DSA, Competitive Programming, JavaScript, System Design, and many more! If you want to test your coding ability, you may check out the mock test series and participate in the contests hosted on Coding Ninjas!

But suppose you have just started your learning process and are looking for questions from tech giants like Amazon, Microsoft, Uber, etc. In that case, you must look at the problems, interview experiences, and interview bundles for placement preparations.

However, you may consider our paid courses to give your career an edge over others!

Happy Learning!