Custom layer using tensorflow

Question

I am trying to make customized layer which reduces channels by reduce_sum(axis=-1)

For example when input shape is (32,32,128)

I want to change the input shape from (32,32,128) to (None,32,32,128) and

And if the channels have index which will be like [0],[1],[2],[3]………[126],[127] And my customized layer want to do is adding 2or 3 or 4…N channels

Lets say if i want to add only 2 channels Whoch will be [0]+[1],[1]+[2]….[126]+[127] And the output shape will be (32,32,64)

also (None,32,32,64)

For more details lets say i want to add 3 channels which will be [0]+[1]+[2],[3]+[4]+[4]… … [123]+[124]+[125],[126]+[127] And the output shape will be (32,32,44)

Here too (None,32,32,44)

So is it possible to make it?

Is there index in channels? If so it would be kinda easy to make it I think…

Answer 1

x = tf.random.normal(shape=(32,32,128))
    sum_dim = 2
    x = tf.reduce_sum(tf.reshape(x, (x.shape[0], x.shape[1], -1, sum_dim)), -1)
    #[32, 32, 64]

If you want to write the keras layer:

class ChannelSum(keras.layers.Layer):
    def __init__(self, sum_dim=2):
        super(ChannelSum, self).__init__()
        self.sum_dim = sum_dim
    def call(self, x ):
        return tf.reduce_sum(tf.reshape(x, (-1, x.shape[1], x.shape[2], x.shape[3]//self.sum_dim, self.sum_dim)), -1)