Is it possible to add different behavior for training and testing in keras Functional API

Viewed 141

I want to use different behavior in functional API for training and testing. Is it possible? E.g.,

a = Input
b = CONV1(a)
if testing:
    return b
c = CONV2(b)
1 Answers

Yes, this can be achieved by defining custom keras layers.

Example codes:

class diff_behavior_layer(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        self.dense_1=Dense(64)
        super().__init__(**kwargs)
    def call(self, inputs,training=None):
        if training:
          return self.dense_1(inputs)
        else:
          return inputs

inputs = tf.keras.Input(shape=(2,))

x=Dense(64)(inputs)
x=Dense(64)(x)
x=diff_behavior_layer()(x)
outputs=Dense(64)(x)

model = tf.keras.Model(inputs=inputs, outputs=outputs)

model(data,training=True) # flow through 4 dense layer
model(data,training=False) # flow through 3 dense layer

Remark: 'training' must be used as the keyword argument here. You cannot define your own keyword argument like testing, etc..

Related