Custom layers in TensorFlow

By Vanshikha Sharma

In this tutorial, We will learn what custom layers are and how can we implement custom layers in TensorFlow. When you are writing machine learning code you either want to implement codes with layers that are already implemented or implement customized layers according to your problem. Some of the pre-existing layers are dense layer, CONV2D, LSTM etc. Now we will see how to implement custom layers in Tensorflow.

IMPLEMENTATION

The best way to customize your own layer is to extend tf.keras.layers class and implement _init_, build and call functions. _init_ functions are used to initialize the parameters. build method will be used to get the shape of the tensor and can do more initialization. The call method is used for forwarding propagation.

class MyDenseLayer(tf.keras.layers.Layer):
  def __init__(self, num_outputs):
    super(MyDenseLayer, self).__init__()
    self.num_outputs = num_outputs   # _init_ function

  def build(self, input_shape):
    self.kernel = self.add_weight("kernel",
                                  shape=[int(input_shape[-1]),
                                         self.num_outputs])

  def call(self, input):
    return tf.math.add(input, self.kernel)

layer = MyDenseLayer(10)

_ = layer(tf.zeros([5, 5]))

print([var.name for var in layer.trainable_variables])

['my_dense_layer/kernel:0']