Why doesn't this input pass through this simple PyTorch model?
I have an input/tensor whose shape is:
torch.Size([256, 3, 28, 28])

(batch size here is 256, 3 channels, 28x28 image)
And a model like so:
class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Conv2d(3, 28, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(28, 56, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # output: 56 x 16 x 16

            nn.Conv2d(56, 112, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(112, 112, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # output: 112 x 8 x 8

            nn.Conv2d(112, 224, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(224, 224, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # output: 224 x 4 x 4

            nn.Flatten(),
            nn.Linear(224 * 4 * 4, 896),
            nn.ReLU(),
            nn.Linear(896, 512),
            nn.ReLU(),
            nn.Linear(512, 2))

    def forward(self, xb):
        return self.network(xb)

When I attempt to pass data forward, it fails with:
    ...
    return self.network(xb)
  File &quot;/home/stark/anaconda3/envs/torch-env/lib/python3.8/site-packages/torch/nn/modules/module.py&quot;, line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File &quot;/home/stark/anaconda3/envs/torch-env/lib/python3.8/site-packages/torch/nn/modules/container.py&quot;, line 117, in forward
    input = module(input)
  File &quot;/home/stark/anaconda3/envs/torch-env/lib/python3.8/site-packages/torch/nn/modules/module.py&quot;, line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File &quot;/home/stark/anaconda3/envs/torch-env/lib/python3.8/site-packages/torch/nn/modules/linear.py&quot;, line 93, in forward
    return F.linear(input, self.weight, self.bias)
  File &quot;/home/stark/anaconda3/envs/torch-env/lib/python3.8/site-packages/torch/nn/functional.py&quot;, line 1690, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: mat1 dim 1 must match mat2 dim 0

What am I missing?
Thanks!


nn.MaxPool2d(2, 2),  # output: 56 x 16 x 16

This is wrong. The original input is of size (256, 3, 28, 28). The convolutional and ReLU layers you're using do not change the batch, height, or width dimensions; they only change the &quot;channel&quot; dimension. Prior to the max pooling layer, the tensor size is (256, 56, 28, 28). The max pooling layer has a kernel size of 2 and a stride of 2, so it will cut both the height and width in half. So the output of this max pooling layer is of size (256, 56, 14, 14).
For the same reason, the output of the next max pooling layer is of size (256, 112, 7, 7), and the output of the last max pooling layer is of size (256, 224, 3, 3).
So you can either fix this by changing your input size to (256, 3, 32, 32), if that's an option, or changing the first linear layer to nn.Linear(224 * 3 * 3, 896)
