Tensorflow: Custom Layer Output Naming

user_6920649 0
Viewed 97

I'm using a tf.data pipeline with dicts to train my model and therefore the input and output names of my model have to match the dict-keys in my tf.data.Dataset. To archive this, I need a way to control the output names of my custom layers.

I can not figure out how to do this. See the example below:

import tensorflow as tf

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, name=None):
        super(CustomLayer, self).__init__(name=name)
        self.dense = tf.keras.layers.Dense(32)

    def __call__(self, inputs):
        x = self.dense(inputs)
        x = tf.add(x, 42)
        return x

m = tf.keras.models.Sequential()
m.add(tf.keras.Input(shape=(100,)))
m.add(tf.keras.layers.Dense(84))
m.add(CustomLayer(name='custom_layer'))

m.summary()

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense (Dense)                (None, 84)                8484      
_________________________________________________________________
dense_1 (Dense)              (None, 32)                2720      
_________________________________________________________________
tf.math.add (TFOpLambda)     (None, 32)                0         
=================================================================
Total params: 11,204
Trainable params: 11,204
Non-trainable params: 0

What I expect or want to achieve is a layer naming like the following:

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense (Dense)                (None, 84)                8484      
_________________________________________________________________
custom_layer                 (None, 32)                2720             
=================================================================
Total params: 11,204
Trainable params: 11,204
Non-trainable params: 0
tensorflow2.0 tensorflow python
1 Answers

Well, it turns out the issue was that I was overwriting __call__ instead of call. The following solves the above issue:

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, name):
        super(CustomLayer, self).__init__(name=name)
        self.dense = tf.keras.layers.Dense(32, name='abc')

    def call(self, inputs):
        x = self.dense(inputs)
        x = x * 2
        return x

m = tf.keras.models.Sequential()
m.add(tf.keras.Input(shape=(100,)))
m.add(tf.keras.layers.Dense(84))
m.add(CustomLayer(name='some_name'))

m.summary()

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense (Dense)                (None, 84)                8484      
_________________________________________________________________
some_name (CustomLayer)      (None, 32)                2720             
=================================================================
Total params: 11,204
Trainable params: 11,204
Non-trainable params: 0
user_69206490
Related