Hi all,
I need to implement below model structure defined in Keras using pytorch.
    # MNIST model
    layers = [
        Conv2D(64, (3, 3), padding='valid', input_shape=(28, 28, 1)),
        Activation('relu'),
        Conv2D(64, (3, 3)),
        Activation('relu'),
        MaxPooling2D(pool_size=(2, 2)),
        Dropout(0.5),
        Flatten(),
        Dense(128),
        Activation('relu'),
        Dropout(0.5),
        Dense(10),
        Activation('softmax')
    ]

I tried below model in Pytorch, but I would like to double check if I did it correctly:
class LeNet_dropout(nn.Module):
    def __init__(self):
        super(LeNet_dropout,self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, padding=0, bias=True)
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=0, bias=True)
        self.dropout = nn.Dropout(0.5)
        self.fc1 = nn.Linear(12*12*64, 128)
        self.fc2 = nn.Linear(128, 10)


    def last_hidden_layer_output(self, X):
        X = F.relu(self.conv1(X))
        X = self.dropout(F.max_pool2d(F.relu(self.conv2(X)),2))
        X = X.view(-1, 12*12*64)
        X = self.dropout(F.relu(self.fc1(X)))
        return X

    def forward(self, X):
        X = F.relu(self.conv1(X))
        X = self.dropout(F.max_pool2d(F.relu(self.conv2(X)),2))
        X = X.view(-1, 12*12*64)
        X = self.dropout(F.relu(self.fc1(X)))
        X = torch.softmax(F.relu(self.fc2(X)), dim=1)
        return X

Would you please comment and correct me if I did anything wrong?
And I am not sure which loss function to use to train this model. Because there is a softmax in output. I would appreciate if you can also guide how to train exactly? I think I cannot use nn.CrossEntropyLoss as below:
output = Model(data)
loss = nn.CrossEntropyLoss(output,target)
And last thing, I will also need to use the output of to last hidden layer activation. Thats why I also include a method called last_hidden_layer_output. Is there any other way to handle this?