Coupling of Different Blocks in a UNET - keras

Question

I am starting to work with Neuralnetworks using Keras. I try to adapt the model (UNet-like architecture) given by Sim, Oh, Kim, Jung in "Optimal Transport driven CycleGAN for Unsupervised Learning in Inverse Problems" (Fig. 10).
def def_generator(image_shape=(256,256,3)):
init= RandomNormal(stddev=0.02)
#Start 1st Block
in_image = Input(shape=image_shape)
g1=Conv2D(64,(3,3))(in_image)
g1=InstanceNormalization(axis=-1)(g1)
g1=LeakyReLU(alpha=0.2)(g1)
g1=Conv2D(64,(3,3))(g1)
g1=InstanceNormalization(axis=-1)(g1)
g1=LeakyReLU(alpha=0.2)(g1)
#End of 1st Block
#Start of 2nd Block
g2=MaxPool2D()(g1)
g2=Conv2D(128,(3,3))(g2)
g2=InstanceNormalization(axis=-1)(g2)
g2=LeakyReLU(alpha=0.2)(g2)
g2=Conv2D(128,(3,3))(g2)
g2=InstanceNormalization(axis=-1)(g2)
g2=LeakyReLU(alpha=0.2)(g2)
#End of 2nd Block
#Start of 3rd Block
g3=MaxPool2D()(g2)
g3=Conv2D(256,(3,3))(g3)
g3=InstanceNormalization(axis=-1)(g3)
g3=LeakyReLU(alpha=0.2)(g3)
g3=Conv2D(256,(3,3))(g3)
g3=InstanceNormalization(axis=-1)(g3)
g3=LeakyReLU(alpha=0.2)(g3)
#End of 3rd Block
#Start of 4th block
g4=MaxPool2D()(g3)
g4=Conv2D(512,(3,3))(g4)
g4=InstanceNormalization(axis=-1)(g4)
g4=LeakyReLU(alpha=0.2)(g4)
g4=Conv2D(512,(3,3))(g4)
g4=InstanceNormalization(axis=-1)(g4)
g4=LeakyReLU(alpha=0.2)(g4)
g4=Conv2D(256,(3,3))(g4)
g4=InstanceNormalization(axis=-1)(g4)
g4=LeakyReLU(alpha=0.2)(g4)
g4=Conv2DTranspose(256,(2,2),strides=(4,4),output_padding=1)(g4)
#End of 4th Block
#Start of 5th Block
g5input=Concatenate()([g4,g3])
g5=Conv2D(256,(3,3))(g5input)
g5=InstanceNormalization(axis=-1)(g5)
g5=LeakyReLU(alpha=0.2)(g5)
g5=Conv2D(256,(3,3))(g5)
g5=InstanceNormalization(axis=-1)(g5)
g5=LeakyReLU(alpha=0.2)(g5)
g5=Conv2DTranspose(128,(2,2),strides=(3,3), padding='same', output_padding=0)(g5)
#End of 5th Block
#Start of 6th block
g6input=Concatenate()([g5,g2])
g6=Conv2D(128,(2,2))(g6input)
g6=InstanceNormalization(axis=-1)(g6)
g6=LeakyReLU(alpha=0.2)(g6)
g6=Conv2D(128,(2,2))(g6)
g6=InstanceNormalization(axis=-1)(g6)
g6=LeakyReLU(alpha=0.2)(g6)
g6=Conv2DTranspose(64,(2,2),strides=(2,2), padding='valid', output_padding=1)(g6)
#End of 6th Block
#Start of 7th block
g7input=Concatenate()([g6,g1])
g7=Conv2D(64,(2,2))(g7input)
g7=InstanceNormalization(axis=-1)(g7)
g7=LeakyReLU(alpha=0.2)(g7)
g7=Conv2D(64,(2,2))(g7)
g7=InstanceNormalization(axis=-1)(g7)
g7=LeakyReLU(alpha=0.2)(g7)
g7=Conv2DTranspose(1,(1,1))(g7)
model=Model(in_image, g5)
model.compile(loss='mse', optimizer=Adam(lr=2e-4,beta_1=0.5), loss_weights=[0.5], metrics=['accuracy'])
return model
g=def_generator((120,120,1))
print(g.summary())
I run always in the problem that the dimensions of the layers which should be concatenated are not compatible.
I understand that this issue is resulting from the MaxPooling+Conv2d steps before.
I am now wondering if there is a trick/strategy to avoid/reduce this issue?
Any help will be appreciated.
Best wishes
Michael

Answer 1

the problem is very simple, you are concatenating block with layers with different size, this is happening because you are trying to run the network on images that are NOT POWER OF 2 size, when you do the max pooling of an image that is not divisible for 2 you lose a pixel (243x243 -> 121x121) and when you double with the traspose you get a different size (121x121 -> 242x242) and the concatenation doesnt work because 242 is different to 243, the images are of different size (at least this is what i think, you should have shared the error).
This means that when an image reaches a maxpooling layer it needs to have an edge divisible for 2.
so, solution:
having 4 blocks means that the images need to be AT LEAST divisible for 16, otherwise it will not work