I’m doing some tinkering with a modified AlexNet and adding in some BatchNorm to look at the position of batchnorm in relation to the activation function, and I’m getting a dimensions error, and I can’t seem to figure out where it’s coming from. I more or less copied the AlexNet architecture from the PyTorch code 2, but added in BatchNorm. Here’s my model class, train function, and training loop:
 class AlexNet_bn_first(nn.Module):
    def __init__(self):
        super(AlexNet_bn_first, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.BatchNorm2d(192),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.BatchNorm2d(384),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.BatchNorm2d(4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.BatchNorm2d(4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, 6),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        
        return F.log_softmax(x, dim=1)

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()

for epoch in range(num_epochs):
    train(model, device, train_loader, optimizer, epoch)
    test_loss = test(model, device, test_loader)
    is_best = test_loss < best_loss
    test_losses.append(test_loss)
    save_checkpoint({
        'epoch': epoch,
        'state_dict': model.state_dict(),
        'best_loss': best_loss,
        'optimizer': optimizer.state_dict(),
    }, is_best)

And here’s the error traceback:
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-16-495c1ea3f5ba> in <module>
      5 
      6 for epoch in range(1, num_epochs+1):
----> 7     train(model, device, train_loader, optimizer, epoch)
      8     test_loss = test(model, device, test_loader)
      9     is_best = test_loss < best_loss

<ipython-input-3-5c27374f04c3> in train(model, device, train_loader, optimizer, epoch, log_interval)
      4         data, target = data.to(device), target.to(device)
      5         optimizer.zero_grad()
----> 6         output = model(data)
      7         loss = F.nll_loss(output, target)
      8         loss.backward()

~\AppData\Local\Continuum\anaconda3\envs\detection\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    491             result = self._slow_forward(*input, **kwargs)
    492         else:
--> 493             result = self.forward(*input, **kwargs)
    494         for hook in self._forward_hooks.values():
    495             hook_result = hook(self, input, result)

<ipython-input-12-199dfe1fc3cc> in forward(self, x)
     39         x = self.avgpool(x)
     40         x = torch.flatten(x, 1)
---> 41         x = self.classifier(x)
     42 
     43         return F.log_softmax(x, dim=1)

~\AppData\Local\Continuum\anaconda3\envs\detection\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    491             result = self._slow_forward(*input, **kwargs)
    492         else:
--> 493             result = self.forward(*input, **kwargs)
    494         for hook in self._forward_hooks.values():
    495             hook_result = hook(self, input, result)

~\AppData\Local\Continuum\anaconda3\envs\detection\lib\site-packages\torch\nn\modules\container.py in forward(self, input)
     90     def forward(self, input):
     91         for module in self._modules.values():
---> 92             input = module(input)
     93         return input
     94 

~\AppData\Local\Continuum\anaconda3\envs\detection\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    491             result = self._slow_forward(*input, **kwargs)
    492         else:
--> 493             result = self.forward(*input, **kwargs)
    494         for hook in self._forward_hooks.values():
    495             hook_result = hook(self, input, result)

~\AppData\Local\Continuum\anaconda3\envs\detection\lib\site-packages\torch\nn\modules\batchnorm.py in forward(self, input)
     59     @weak_script_method
     60     def forward(self, input):
---> 61         self._check_input_dim(input)
     62 
     63         # exponential_average_factor is self.momentum set to

~\AppData\Local\Continuum\anaconda3\envs\detection\lib\site-packages\torch\nn\modules\batchnorm.py in _check_input_dim(self, input)
    248         if input.dim() != 4:
    249             raise ValueError('expected 4D input (got {}D input)'
--> 250                              .format(input.dim()))
    251 
    252 

ValueError: expected 4D input (got 2D input)

It looks like the error is happening with the input to self.classifier. Indeed, I do flatten the tensor to 2D before it goes into the classifier, and when I check model.classifier, this is what I get:
Sequential(
  (0): Dropout(p=0.5)
  (1): Linear(in_features=9216, out_features=4096, bias=True)
  (2): BatchNorm2d(4096, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (3): ReLU(inplace)
  (4): Dropout(p=0.5)
  (5): Linear(in_features=4096, out_features=4096, bias=True)
  (6): BatchNorm2d(4096, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (7): ReLU(inplace)
  (8): Linear(in_features=4096, out_features=6, bias=True)
)

So, it should be expecting a 2D input, no? What am I missing here?