Build a Flatten layer to work with layers that have variable input shape

Question

Let's say I have this model:

model = Sequential()
model.add(Conv2D(64, (3,3), input_shape=(None, None, 1), strides=1, padding='same'))
model.add(Flatten())
model.add(Dense(2, activation='softmax'))
model.summary()

It will generate the following error:

ValueError: The last dimension of the inputs to `Dense` should be defined. Found `None`.

I want to make a new layer that will be able to build a flatten tensor from tensors with variable shape (None, None, 1).

I tried to build a new layer like the following:

class MyFlattenLayer(tf.keras.layers.Layer):
    def __init__(self, data_format=None, **kwargs):
        super(MyFlattenLayer, self).__init__(**kwargs)
        self.data_format = conv_utils.normalize_data_format(data_format)
        self.input_spec = InputSpec(min_ndim=1)
        self._channels_first = self.data_format == 'channels_first'
        

    def call(self, inputs):
        shape = K.shape(inputs)
        batchSize = shape[:1]
        newShape = K.variable([-1],dtype='int32')
        newShape = K.concatenate([batchSize,newShape])
        return K.reshape(inputs,newShape)
    
    def compute_output_shape(self, input_shape):
        batchSize = input_shape[:1]
        newShape = K.variable([-1],dtype='int32')
        newShape = K.concatenate([batchSize,newShape])
        
        return tensor_shape.TensorShape(newShape)

    def get_config(self):
        config = super(MyFlattenLayer, self).get_config()
        config.update({'data_format': self.data_format})
        return config

model = Sequential()
model.add(Conv2D(64, (3,3), input_shape=(None, None, 1), strides=1, padding='same'))
model.add(MyFlattenLayer())
model.add(Dense(2, activation='softmax'))
model.summary()

But this also generates another error:

ValueError: The last dimension of the inputs to `Dense` should be defined. Found `None`.

How to build a Flatten layer to work with layers that have variable input shape ?

UPDATE

I was able to build a MyFlattenLayer for variable dimension of input for first Conv2D layer. It compiles without errors, and I can even train and predict just for testing. But something looks strange, like when I try to show summary of the model and count parameters, and even when I try to display all tensors values for all layers, then values of MyFlattenLayer are missing...and not displayed, however it is able to use compile, train and predict functions, How if it is missing some parameters? I understand that Dense layer must have a constant dimension of input, but still looking to find a workaround or something that will help to achive something similar not necessarily the same.

import tensorflow.keras.backend as K
from tensorflow.python.keras.utils import conv_utils
from tensorflow.python.keras.engine.input_spec import InputSpec
from tensorflow.python.framework import tensor_shape

from tensorflow.python.keras import backend
from tensorflow.python.ops import array_ops


class MyFlattenLayer(Flatten):

  def __init__(self, data_format=None, **kwargs):
    super(MyFlattenLayer, self).__init__(**kwargs)
    self.data_format = conv_utils.normalize_data_format(data_format)
    self.input_spec = InputSpec(min_ndim=1)

  def compute_output_shape(self, input_shape):
    input_shape = tensor_shape.TensorShape(input_shape).as_list()
    if self.data_format == 'channels_last':
      return tensor_shape.TensorShape([input_shape[0], input_shape[3]])
    else:
      return tensor_shape.TensorShape([input_shape[0], input_shape[1]])

  def call(self, inputs):
    tens = backend.flatten(inputs)
    threshold = tf.constant(-99999999, dtype=tf.float32)
    y = tf.cast(inputs < threshold, dtype=tf.float32)
    if tens.shape[0] == None:
      tens = array_ops.reshape(inputs, (y.shape[3],1)) 
      return tens
    else:
      tens = array_ops.reshape(inputs, (tens.shape[0],1)) #backend.reshape(tens, (tens.shape[0],1)) 
      
      return tens
    

  def get_config(self):
    config = {'data_format': self.data_format}
    base_config = super(MyFlattenLayer, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

model = Sequential()
model.add(Conv2D(64, (3,3), input_shape=(None, None, 1), strides=1, padding='same'))
model.add( MyFlattenLayer() )
# model.add( Flatten() )
model.add(Dense(2, activation='softmax'))
model.summary()
optimizer = Adam(lr = 0.00006)
model.compile(loss="binary_crossentropy",optimizer=optimizer)
x = np.arange(0, 25).reshape(1, 5, 5, 1).astype(np.float32)
out = model.train_on_batch(x, np.array([[1, 0]]))
print("results: \n", out , "\n\n")
for layer in model.layers: print(layer.get_config(), layer.get_weights())

This also works:

x = np.arange(0, 25*35).reshape(1, 25, 35, 1).astype(np.float32)
out = model.train_on_batch(x, np.array([[1, 0]]))

Answer 1

The problem is not that the Flatten operation does not work, it theoretically should work for variable shapes as far as I am concerned. The problem is, that you try to feed the flattened output into a dense layer, which needs to know how many weights it has. The number of weights does directly depend on the input shape (more precisely on the last dimension of it, since the dense layer is applied to this dimension), so feeding a variable shaped input into the dense layer does not work.