How to reshape last layer of pytorch CNN model while doing transfer learning

Actually i am trying to replicate keras structure to pytorch.(new in pytorch).
Here is keras architecture

base_model = InceptionV3(weights='imagenet', include_top=False)
x = base_model.output
x = Dense(512, activation='relu')(x)
predictions = Dense(49*6,activation='sigmoid')(x)
reshape=Reshape((49,6))(predictions)
model = Model(inputs=base_model.input, outputs=reshape)
for layer in base_model.layers:
    layer.trainable = False

I want to reshape last layer of my netwrok. I have implemented transfer learning.

model =  models.inception_v3(pretrained=True)

for param in model.parameters():
    param.requires_grad = False
    
num_ftrs = model.fc.in_features

I believe if i can attach last layer of resnet with my following architecture, problem can be solved. But i dont know how i can attach them

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.fc1 = nn.Linear(num_ftrs, 512)
        self.fc2 = nn.Linear(512, 49*6)

    def forward(self, x):
        print (x.shape)
        x = x.view(-1,num_ftrs)
        #print (x.shape)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        x=torch.sigmoid(x.view(49,6))
        return x

Any idea, how this problem can be resolved

You could add your custom linear layers to model.fc as shown here:

model =  models.inception_v3()
fc_feat = model.fc.in_features
model.fc = nn.Sequential(
    nn.Linear(fc_feat, 512),
    nn.ReLU(),
    nn.Linear(512, 49*6),
    nn.Sigmoid(),
)

Let me know, if that would work for your use case.

there is no reshaping of last layer output, i want 49,6 instead of 49*6

You can reshape the output tensor of shape (49*6, ) to (49, 6).
Simply do:

output = model(input)
output = output.view(-1, 49, 6)
2 Likes