Add a Truncated layer in Functional API model

Question

I am trying to add truncation layer after Conv2D layer in the following code:

input_layer = Input(shape=(256, 256, 1))
conv = Conv2D(8, (5, 5), padding='same', strides=1, use_bias=False)(input_layer)
output_layer = Activation(activation='tanh')(lambda_layer)
output_layer = AveragePooling2D(pool_size= (5, 5), strides=2)(output_layer)
output_layer = BatchNormalization()(output_layer)

The truncation layer must satisfy:

−T if x < −T
x if −T ≤ x ≤ T
T  x > T 

where `T` is a threshold value, x= the output of convolution layer`

Could someone please help me to build this layer?

Thank you

Answer 1

you can build the desired function with tensorflow.keras.backed.switch and wrap it inside a Lambda layer

build and test the function:

T = 5
X = tf.constant(np.random.uniform(-10, 10, (3,5)))

def switch_func(X, T):
    
    zeros = tf.zeros_like(X)
    T_matrix = tf.ones_like(X) * T

    cond1 = K.switch(X < -T_matrix, -T_matrix, zeros)
    cond2 = K.switch(X > T_matrix, T_matrix, zeros)
    cond3 = K.switch(tf.abs(cond1 + cond2) == T, zeros, X)
    res = cond1 + cond2 + cond3
    return res

switch_func(X, T)

<tf.Tensor: shape=(3, 5), dtype=float64, numpy=
array([[-5.        ,  0.65807168, -4.93481499, -5.        , -2.94954848],
       [-1.25114075, -5.        ,  2.97657545,  5.        , -0.8958152 ],
       [-1.26611956,  5.        , -3.38477137,  5.        , -3.53358454]])>

usage inside the model:

X = np.random.uniform(0,1, (100,10))
y = np.random.uniform(0,1, (100,))

inp = Input((10,))
x = Dense(8)(inp)
x = Lambda(lambda x: switch_func(x, T=0.5))(x)
out = Dense(1)(x)

model = Model(inp, out)
model.compile('adam', 'mse')
model.fit(X,y, epochs=3)