Reusable block in Keras' functional API

Question

The goal is to create a block of layers, using Keras' functional API, which is usable (also syntax-wise) like a 'normal' Keras layer. Here is a toy example

from tensorflow.keras import layers as kl

def layer_block(prev_layer, args):
    # some code using 'args'
    layer = kl.Dense(units=prev_layer.shape[1])(prev_layer)
    layer = kl.Dense(units=5)(layer)
    layer = kl.Dense(units=prev_layer.shape[1])(layer)

    return layer

This block is called using layer_block(prev_layer, args) which is in contradiction to Keras' functional API's syntax. It should rather look like layer_block(args)(prev_layer).

The approach so far is to wrap this block by another block:

def outer_block(args):
    def layer_block(prev_layer, args):
        # some code using 'args'
        layer = kl.Dense(units=prev_layer.shape[1])(prev_layer)
        layer = kl.Dense(units=5)(layer)
        layer = kl.Dense(units=prev_layer.shape[1])(layer)

        return layer
    return lambda prev_layer: layer_block(prev_layer, args)

Now two questions arise:

Is there an easier way to achieve this?
Is it effective this way or does it have negative impact on performance?

Thank you in advance!

user_2097240 · Answer

What you're doing doesn't affect performance, you're creating layers perfectly fine.

There is no problem in any of your two approaches, but if you do want to make it work as an actual layer, transform it into a model.

This may not work in every keras version:

class LayerBlock(tensorflow.keras.Model): #not sure if it works in normal keras (without tf)

    def __init__(self):
        super(LayerBlock, self).__init__(outer_units)
        self.layer1 = kl.Dense(units=outer_units)
        self.layer2 = kl.Dense(units=5)
        self.layer3 = kl.Dense(units=outer_units)

    def call(self, inputs):
        x = self.layer1(inputs)
        x = self.layer2(x)
        x = self.layer3(x)
        return x

This tutorial seems to suggest that you can use tf.keras.Layer instead of tf.keras.Model, but that sounds strange to me. It may work with eager mode on, but it lacks the build method with a self.built=True statement.

Reusable block in Keras' functional API

1 Answers