Lazy Man's Guide to Convolution Size (PyTorch)

In Keras, one can add a linear layer after a conv layer, and the library will automatically determine the ‘size’ for the input dimension.

Not in PyTorch!

In PyTorch, we have to manually calculate (using a quirky formula mentioned in the documentation and other places). Here is a quick fix!

First, specify and build the model without the linear layer. e.g for Fashion MNIST

class CNN(nn.Module):
    def __init__(self, N):
        super(CNN, self).__init__()
        self.cl = nn.Sequential(
                                nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=2),
                                nn.ReLU(),
                                nn.Conv2d(in_channels=16, out_channels=128, kernel_size=3, stride=1),
                                nn.ReLU(),
                                nn.Conv2d(in_channels=128, out_channels=32, kernel_size=3, stride=2),
                                nn.ReLU()
                                )

    
    def forward(self, X):
        out = self.cl(X)
        return out

Then run the following code

model = CNN(10)

from torchsummary import summary
summary(model, (1,28,28))

This will print the shape of the output for each layer. Just note the size of the last layer. Here, it will be [-1, 32, 5, 5].

Now, your next Dense layers must have the number of neurons as 32 * 5 * 5 = 800!