Pytorch 矩阵大小问题不会成倍增加

Question

This may have been answered before so happy about any links. I am new to pytorch and do not understand why my Conv2d pipeline is failing with

mat1 and mat2 shapes cannot be multiplied (64x49 and 3136x512)

        self.net = nn.Sequential(
            nn.Conv2d(in_channels=c, out_channels=32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2),
            nn.ReLU(),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(3136, 512),
            nn.ReLU(),
            nn.Linear(512, output_size)
        )

with input shape 1x84x84 .

I did the calculation and this is the size that breaks down over the different steps with the kernel and size settings per layer.

84 ->  K:8 , S:4 => 20
20 -> K:3 , S:2 => 9
9 -> K:3 , S:1 => 7

7^2 * 64 => 3136  for the flattened layer

I am not sure where the 64x49 is coming from .

Answer 1

I have tried your model and your calculation is totally correct. The problem lies in your input. For torch calculation, if your input shape is 1x84x84, a 3d torch, you should input a 4d torch indeed, where the first dimension represent the batch-size. You may find more information about batch, which is widely used to enhance computation speed.

If you just want to test on single data, you may just add a dimension like x = x[None, :] to make it become 4d torch. This will be a quick fix to your problem