Translate CNN from keras to pytorch - image

Question

I need some help. I am trying to convert a CNN from keras to pytorch. I am reconstructing an MR image. The input is the image coming from the scanner in the fourier domain and the output is the reconstructed image. The input image has two channels (first channel: real part, second channel: imaginary part). Unfortunately, the results are quite different, so I believe I am doing something wrong. I just cannot find out myself what it is. Here is the keras code:
def AUTOMAP_Basic_Model(param):
fc_1 = keras.Input(shape=(64,64,2), name='input')
fc_2 = layers.Conv2D(64,(64,1),strides=1,padding='same',activation='relu')(fc_1)
fc_4 = layers.Conv2D(64,(1,64),strides=1,padding='same',activation='relu')(fc_2)
fc_4 = layers.Conv2D(64,(64,1),strides=1,padding='same',activation='relu')(fc_4)
fc_5 = layers.Conv2D(64,(1,64),strides=1,padding='same',activation='relu')(fc_4)
c_1 = layers.Conv2D(64,5,strides=1,padding='same',activation='relu')(fc_5)
c_2 = layers.Conv2D(64,5,strides=1,padding='same',activation='relu')(c_1)
c_3 = layers.Conv2DTranspose(1,7,strides=1,activation='sigmoid',padding='same')(c_2)
model = keras.Model(inputs = fc_1,outputs = c_3)
return model
And this is my translation to pytorch:
class AUTOMAP_Basic_Model(nn.Module):
def __init__(self, inputShape, nrFilters):
super(AUTOMAP_Basic_Model, self).__init__()
self.conv1 = nn.Conv2d(2, 64, (64,1), padding='same')
self.relu = nn.ReLU()
self.conv2 = nn.Conv2d(64, 64, (1,64), padding ='same')
self.conv3 = nn.Conv2d(64, 64, (64,1), padding='same')
self.conv4 = nn.Conv2d(64,64,5,padding='same')
self.conv5 = nn.Conv2d(64,64,5,padding='same')
self.convTranspose = nn.ConvTranspose2d(64,1,7,padding=3,output_padding=0)
self.sigmoid = nn.Sigmoid()
self.tan = nn.Tanh()
def forward(self, x):
batch_size = len(x)
out = self.conv1(x)
out = self.relu(out)
out = self.conv2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.relu(out)
out = self.conv2(out)
out = self.relu(out)
out = self.conv4(out)
out = self.relu(out)
out = self.conv5(out)
out = self.relu(out)
out = self.convTranspose(out)
out=self.sigmoid(out)
return out
I am new to pytorch, thats why I do not know what could be wrong. While the keras model is converging to the right image, the pytorch model is giving me a constant value of 0.45 for every pixel in the image.

Answer 1

It seems like you have some strange ordering and naming in both Keras and PyTorch. In the first one fc_3 is missing, which might still be viable but I don't understand why.
In the PyTorch implementation, you need to have 4 conv layer (fc_2, fc_3, fc_4, fc_5), but you only implemented 3 (conv1, conv2, conv3).
I'd rewrite the model with something like:
class AUTOMAP_Basic_Model(nn.Module):
def __init__(self, inputShape, nrFilters):
super(AUTOMAP_Basic_Model, self).__init__()
self.conv1 = nn.Conv2d(2, 64, (64, 1), padding='same')
self.conv2 = nn.Conv2d(64, 64, (1, 64), padding='same')
self.conv3 = nn.Conv2d(64, 64, (64, 1), padding='same')
self.conv4 = nn.Conv2d(64, 64, (1, 64), padding='same')
self.conv5 = nn.Conv2d(64, 64, 5, padding='same')
self.conv6 = nn.Conv2d(64, 64, 5, padding='same')
self.relu = nn.ReLU(inplace=True)
self.convTranspose = nn.ConvTranspose2d(64, 1, 7, padding=3, output_padding=0)
self.sigmoid = nn.Sigmoid()
self.tan = nn.Tanh()
def forward(self, x):
batch_size = len(x)
out = self.relu(self.conv1(x))
out = self.relu(self.conv2(out))
out = self.relu(self.conv3(out))
out = self.relu(self.conv4(out))
out = self.relu(self.conv5(out))
out = self.relu(self.conv6(out))
out = self.convTranspose(out)
out = self.sigmoid(out)
return out
I'd also suggest using Conv+BN+Relu