Changing Keras Layer Output Shape

Question

I am working on a program that replaces strided Conv2D layers (strides = 2) with an unstrided Conv2D layer (strides = 1), and then adds an AveragePooling2D layer (strides = 2) after the activation layer. In other words, instead of having the Conv2D layer reduce the output dimensions, the AveragePooling2D layer will do that.

I use the methodology described in this post to replace the Conv2D layer with the unstrided version, and to insert the AveragePooling2D layer after the activation. Replacing the Conv2D layer works perfectly. The trouble is that the new output of the activation layer will have a different shape than it used to. This post explains that the "multiple" output shape is because there are now two output nodes attached to the activation layer, one with the old shape from the strided convolution, and one with the new shape from the unstrided convolution.

Here is code to illustrate the problem:

# Origninal model with strided convolution
inp = tf.keras.layers.Input((32,32,3))
x = tf.keras.layers.Conv2D(filters=1,kernel_size=3,strides=2, padding='same')(inp)
x = tf.nn.relu(x)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dense(units=10, activation='linear')(x)
outp = x
model = tf.keras.models.Model(inp,outp)

# Clone original model in order to modify the clone
model_clone = tf.keras.models.clone_model(model)

# Customize the cloned model - this function is where the convolution
# is replaced and the average pooling layer is added
model_custom = MyCustomFunction(model_clone)

Summary of original model:

Model: "model_2"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         [(None, 32, 32, 3)]       0         
_________________________________________________________________
conv2d_2 (Conv2D)            (None, 16, 16, 1)         28        
_________________________________________________________________
tf.nn.relu_1 (TFOpLambda)    (None, 16, 16, 1)         0         
_________________________________________________________________
global_average_pooling2d_1 ( (None, 1)                 0         
_________________________________________________________________
dense_1 (Dense)              (None, 10)                20        
=================================================================
Total params: 48
Trainable params: 48
Non-trainable params: 0
_________________________________________________________________

Summary of model_custom:

Model: "model_6"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_3 (InputLayer)         [(None, 32, 32, 3)]       0         
_________________________________________________________________
conv2d_6 (Conv2D)            (None, 32, 32, 1)         28        
_________________________________________________________________
tf.nn.relu_2 (TFOpLambda)    multiple                  0         
_________________________________________________________________
average_pooling2d (AveragePo (None, 16, 16, 1)         0         
_________________________________________________________________
global_average_pooling2d_2 ( (None, 1)                 0         
_________________________________________________________________
dense_2 (Dense)              (None, 10)                20        
=================================================================
Total params: 48
Trainable params: 48
Non-trainable params: 0
_________________________________________________________________

When I try to train the new model_custom, I receive this error:

InvalidArgumentError:  logits and labels must have the same first dimension, got logits shape [32768,1] and labels shape [32]

I believe this is because it is still looking at the old output shape of the activation layer (None, 16, 16, 1) rather than the new one (None, 32, 32, 1).

Here is the code I use for training, for reference:

model_custom.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

history = model_custom.fit(train_images, train_labels, epochs=2,
                    validation_data=(test_images, test_labels))

When I apply the same code to the original model, it works fine.

Any help would be greatly appreciated!

Answer 1

No, the error InvalidArgumentError: logits and labels must have the same first dimension, got logits shape [32768,1] and labels shape [32] is not related to model architecture, but rather because of the batch size dimension. Refer to this .

You fed the model data of shape [32768,1] not batched, while the loss you applied wants a shape of [32] , to one hot encode on its own.