How to write a new custom layer in Keras?

Question

I'm currently engaged in a project where I need to create a new layer in Keras. This new layer should be able to rearrange the input tensor. I've successfully implemented the custom layer with code that functions as intended. However, when I integrate this custom layer into a model structure, I encounter an error related to dimensions.

This is the custom layer code:

# Coustom Layer
import tensorflow as tf
from tensorflow.keras.layers import Layer
import numpy as np

class RearrangeLayer(Layer):
    def call(self, inputs):
        num_rows, num_columns = tf.shape(inputs)[1], tf.shape(inputs)[2]

        # Create indices for gathering columns
        indices = tf.range(num_columns)

        # Gather even and odd columns
        even_columns = tf.gather(inputs, indices[::2], axis=2)
        even_columns = tf.reshape(even_columns, (tf.shape(inputs)[0], -1, 2 * even_columns.shape[-1]))
        odd_columns = tf.gather(inputs, indices[1::2], axis=2)
        odd_columns = tf.reshape(odd_columns, (tf.shape(inputs)[0], -1, 2 * odd_columns.shape[-1]))

        # Interleave the rows alternatively
        output_data = tf.zeros((1, even_columns.shape[1] + odd_columns.shape[1], even_columns.shape[2]), dtype='float32')
        output_data = np.array(output_data)
        output_data[0, 0::2] = even_columns
        output_data[0, 1::2] = odd_columns

        return output_data

---

To ensure that the layer works as i expect, i entered a tensor example and received the desired rearranged output:

# Example usage
inputs = tf.constant([[
    [1, 2, 3, 4, 5, 6],
    [7, 8, 9, 10, 11, 12],
    [13, 14, 15, 16, 17, 18],
    [19, 20, 21, 22, 23, 24]]
], dtype=tf.float32)

rearrange_layer = RearrangeLayer()
output = rearrange_layer(inputs)

print(output)

[[[ 1.  3.  5.  7.  9. 11.]
  [ 2.  4.  6.  8. 10. 12.]
  [13. 15. 17. 19. 21. 23.]
  [14. 16. 18. 20. 22. 24.]]]

---

Now when I integrate the custom layer in a model architecture, encounter an error...

# The model architecture
from keras.layers import Input, Flatten, Conv1D, Conv2D, GlobalAveragePooling1D, Dense, Reshape
from keras.models import Model

#define input layer
INPUT = Input((100,60))

reshaped_input = Reshape((100, 60, 1))(INPUT)
conv = Conv2D(1, (2,2), (2,2), activation='relu')(reshaped_input)
conv = RearrangeLayer()(conv)
conv = Conv2D(16, (2,2), (2,1), activation='relu')(conv)
conv = Conv2D(32, (2,2), (2,1), activation='relu')(conv)
flatted = Flatten()(conv)
reshaped = Reshape((10752, 1))(flatted)
conv = Conv1D(32, 6, activation='relu')(reshaped)

pooling = GlobalAveragePooling1D()(conv)

output = Dense(1024, activation='relu')(pooling)
output = Dense(1)(output)

model = Model(inputs=INPUT, outputs=output)

# The error 

TypeError                                 Traceback (most recent call last)
Cell In[3], line 9
      7 reshaped_input = Reshape((100, 60, 1))(INPUT)
      8 conv = Conv2D(1, (2,2), (2,2), activation='relu')(reshaped_input)
----> 9 conv = RearrangeLayer()(conv)
     10 conv = Conv2D(16, (2,2), (2,1), activation='relu')(conv)
     11 conv = Conv2D(32, (2,2), (2,1), activation='relu')(conv)

File ~\AppData\Roaming\Python\Python310\site-packages\keras\utils\traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File ~\AppData\Local\Temp\__autograph_generated_filen_2mzai2.py:16, in outer_factory.<locals>.inner_factory.<locals>.tf__call(self, inputs)
     14 odd_columns = ag__.converted_call(ag__.ld(tf).gather, (ag__.ld(inputs), ag__.ld(indices)[1::2]), dict(axis=2), fscope)
     15 odd_columns = ag__.converted_call(ag__.ld(tf).reshape, (ag__.ld(odd_columns), (ag__.converted_call(ag__.ld(tf).shape, (ag__.ld(inputs),), None, fscope)[0], -1, 2 * ag__.ld(odd_columns).shape[-1])), None, fscope)
---> 16 output_data = ag__.converted_call(ag__.ld(tf).zeros, ((1, ag__.ld(even_columns).shape[1] + ag__.ld(odd_columns).shape[1], ag__.ld(even_columns).shape[2]),), dict(dtype='float32'), fscope)
     17 output_data = ag__.converted_call(ag__.ld(np).array, (ag__.ld(output_data),), None, fscope)
     18 ag__.ld(output_data)[0, 0::2] = ag__.ld(even_columns)

TypeError: Exception encountered when calling layer "rearrange_layer_1" (type RearrangeLayer).

in user code:

    File "C:\Users\dpc\AppData\Local\Temp\ipykernel_1928\1750308305.py", line 19, in call  *
        output_data = tf.zeros((1, even_columns.shape[1] + odd_columns.shape[1], even_columns.shape[2]), dtype='float32')

    TypeError: unsupported operand type(s) for +: 'NoneType' and 'NoneType'


Call arguments received by layer "rearrange_layer_1" (type RearrangeLayer):
  • inputs=tf.Tensor(shape=(None, 50, 30, 1), dtype=float32)

Answer 1

The error you are encountering is due to a dimension mismatch in the custom Keras layer. The TypeError: unsupported operand type(s) for +: 'NoneType' and 'NoneType' indicates that the shape of even_columns and odd_columns is not being properly inferred, leading to NoneType dimensions. This causes the subsequent addition operation in the tf.zeros function to fail. To fix this, you need to ensure that the shape of even_columns and odd_columns is properly inferred. You can achieve this by using the tf.shape function to explicitly get the shape of the tensors. Here's the modified call method for your custom layer:

class RearrangeLayer(Layer):
   def call(self, inputs):
       num_rows, num_columns = tf.shape(inputs)[1], tf.shape(inputs)[2]


       # Create indices for gathering columns
       indices = tf.range(num_columns)


       # Gather even and odd columns
       even_columns = tf.gather(inputs, indices[::2], axis=2)
       even_columns = tf.reshape(even_columns, (tf.shape(inputs)[0], -1, 2 * tf.shape(even_columns)[2]))
       odd_columns = tf.gather(inputs, indices[1::2], axis=2)
       odd_columns = tf.reshape(odd_columns, (tf.shape(inputs)[0], -1, 2 * tf.shape(odd_columns)[2]))


       # Interleave the rows alternatively
       rearranged_output = tf.concat([even_columns, odd_columns], axis=1)


       return rearranged_output


from keras.layers import Input, Flatten, Conv1D, Conv2D, GlobalAveragePooling1D, Dense, Reshape
from keras.models import Model


# The model architecture
from keras.layers import Input, Flatten, Conv1D, Conv2D, GlobalAveragePooling1D, Dense, Reshape
from keras.models import Model


#define input layer
INPUT = Input((100,60))


reshaped_input = Reshape((100, 60, 1))(INPUT)
conv = Conv2D(1, (2,2), (2,2), activation='relu')(reshaped_input)
conv = RearrangeLayer()(conv)
conv = Reshape((50, 30, 1))(conv) # Add Reshape layer to reshape output tensor
conv = Conv2D(16, (2,2), (2,1), activation='relu')(conv)
conv = Conv2D(32, (2,2), (2,1), activation='relu')(conv)
flatted = Flatten()(conv)
reshaped = Reshape((10752, 1))(flatted)
conv = Conv1D(32, 6, activation='relu')(reshaped)


pooling = GlobalAveragePooling1D()(conv)


output = Dense(1024, activation='relu')(pooling)
output = Dense(1)(output)


model = Model(inputs=INPUT, outputs=output)