How to make the size of output image the same as the original one to compute loss in CNN?

Question

I am define the CNN model for autoencoder as follows:

filters = (32, 16)
X = Input(shape = (32, 32, 3))

# encode
for f in filters:
    X = Conv2D(filters = f, kernel_size = (3, 3), activation = 'relu')(X)
    X = MaxPooling2D(pool_size = (2, 2), strides = (2, 2), padding = 'same')(X)
    X = BatchNormalization(axis = -1)(X)

# decode
for f in filters[::-1]:
    X = Conv2D(filters = f, kernel_size = (3, 3), activation = 'relu')(X)
    X = UpSampling2D(size = (2, 2))(X)
    X = BatchNormalization(axis = -1)(X)

The model summary is

Model: "functional_13"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_7 (InputLayer)         [(None, 32, 32, 3)]       0         
_________________________________________________________________
conv2d_24 (Conv2D)           (None, 30, 30, 32)        896       
_________________________________________________________________
max_pooling2d_12 (MaxPooling (None, 15, 15, 32)        0         
_________________________________________________________________
batch_normalization_24 (Batc (None, 15, 15, 32)        128       
_________________________________________________________________
conv2d_25 (Conv2D)           (None, 13, 13, 16)        4624      
_________________________________________________________________
max_pooling2d_13 (MaxPooling (None, 7, 7, 16)          0         
_________________________________________________________________
batch_normalization_25 (Batc (None, 7, 7, 16)          64        
_________________________________________________________________
conv2d_26 (Conv2D)           (None, 5, 5, 16)          2320      
_________________________________________________________________
up_sampling2d_12 (UpSampling (None, 10, 10, 16)        0         
_________________________________________________________________
batch_normalization_26 (Batc (None, 10, 10, 16)        64        
_________________________________________________________________
conv2d_27 (Conv2D)           (None, 8, 8, 32)          4640      
_________________________________________________________________
up_sampling2d_13 (UpSampling (None, 16, 16, 32)        0         
_________________________________________________________________
batch_normalization_27 (Batc (None, 16, 16, 32)        128       
=================================================================
Total params: 12,864
Trainable params: 12,672
Non-trainable params: 192
_________________________________________________________________

Because the output image has a different dimensions from the input image, I get the error

InvalidArgumentError:  Incompatible shapes: [128,32,32,3] vs. [128,16,16,32]
     [[node mean_squared_error/SquaredDifference (defined at <ipython-input-7-a9683921f595>:83) ]] [Op:__inference_train_function_21329]

Function call stack:
train_function

and thus can not compute for the loss function. Could you please elaborate on how to solve this issue?

Answer 1

I suggest u use padding='same' in convolutions. Pay attention also to not overwrite your input layer with other variables. you miss also a final output layer with the number of channels equals to the channel of the input image

filters = (32, 16)
inp = Input(shape = (32, 32, 3))

# encode
X = inp
for f in filters:
    X = Conv2D(filters = f, kernel_size = (3, 3), padding = 'same', activation = 'relu')(X)
    X = MaxPooling2D(pool_size = (2, 2), strides = (2, 2), padding = 'same')(X)
    X = BatchNormalization(axis = -1)(X)

# decode
for f in filters[::-1]:
    X = Conv2D(filters = f, kernel_size = (3, 3), padding = 'same', activation = 'relu')(X)
    X = UpSampling2D(size = (2, 2))(X)
    X = BatchNormalization(axis = -1)(X)
out =  Conv2D(filters = 3, kernel_size = (3, 3), padding = 'same')(X)
    
model = Model(inp, out)

The model summary now is

Layer (type)                 Output Shape              Param #   
=================================================================
input_9 (InputLayer)         [(None, 32, 32, 3)]       0         
_________________________________________________________________
conv2d_22 (Conv2D)           (None, 32, 32, 32)        896       
_________________________________________________________________
max_pooling2d_10 (MaxPooling (None, 16, 16, 32)        0         
_________________________________________________________________
batch_normalization_20 (Batc (None, 16, 16, 32)        128       
_________________________________________________________________
conv2d_23 (Conv2D)           (None, 16, 16, 16)        4624      
_________________________________________________________________
max_pooling2d_11 (MaxPooling (None, 8, 8, 16)          0         
_________________________________________________________________
batch_normalization_21 (Batc (None, 8, 8, 16)          64        
_________________________________________________________________
conv2d_24 (Conv2D)           (None, 8, 8, 16)          2320      
_________________________________________________________________
up_sampling2d_10 (UpSampling (None, 16, 16, 16)        0         
_________________________________________________________________
batch_normalization_22 (Batc (None, 16, 16, 16)        64        
_________________________________________________________________
conv2d_25 (Conv2D)           (None, 16, 16, 32)        4640      
_________________________________________________________________
up_sampling2d_11 (UpSampling (None, 32, 32, 32)        0         
_________________________________________________________________
batch_normalization_23 (Batc (None, 32, 32, 32)        128       
_________________________________________________________________
conv2d_26 (Conv2D)           (None, 32, 32, 3)         867       
=================================================================
Total params: 13,731
Trainable params: 13,539
Non-trainable params: 192
_________________________________________________________________