nn.Linear should be mismatch, but it works successfully
I'm confused about the in-feature of nn.linear. For out-feature of the model VGG-19's last nn.MaxPool2d, the tensor size is (512, 7, 7). The model below uses pooling function and resizes the tensor to (512, 49), then uses nn.linear(512, 7) directly. Why can't it work successfully without mismatch problem?

source


'''VGG11/13/16/19 in Pytorch.'''
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable


cfg = {
    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}


class VGG(nn.Module):
    def __init__(self, vgg_name):
        super(VGG, self).__init__()
        self.features = self._make_layers(cfg[vgg_name])
        self.classifier = nn.Linear(512, 7)

    def forward(self, x):
        out = self.features(x)
        out = out.view(out.size(0), -1)
        out = F.dropout(out, p=0.5, training=self.training)
        out = self.classifier(out)
        return out

    def _make_layers(self, cfg):
        layers = []
        in_channels = 3
        for x in cfg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
                           nn.BatchNorm2d(x),
                           nn.ReLU(inplace=True)]
                in_channels = x
        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
        return nn.Sequential(*layers)


Why is the assumption that this code works? I tested it, and got the following shapes, and the expected size mismatch error.

def forward(self, x):
    out = self.features(x) # torch.Size([1, 512, 7, 7])
    out = out.view(out.size(0), -1)  # torch.Size([1, 25088])
    out = F.dropout(out, p=0.5, training=self.training)  # torch.Size([1, 25088])
    out = self.classifier(out)  # RuntimeError: size mismatch, m1: [1 x 25088], m2: [512 x 7]
    return out


One mistake that you made with inferring the sizes, is that you omitted the batch dimension. That is why you may falsely conclude that the shape change with out.view(out.size(0), -1) is [512, 7, 7] -> [512, 49] instead the correct [b, 512, 7, 7] -> [b, 25088] where b is the batch size.

As expected when classifier is changed to 

self.classifier = nn.Linear(25088, 7)


then the forward function works, no size mismatch error.
