Expected 4D tensor as input, got 2D tensor instead
I'm trying to build a neural network using the pre-trained network VGG16 on Pytorch.

I understand that I need to adjust the classifier part of the network, so I have frozen the parameters to prevent backpropagation through them.

Code:

%matplotlib inline
%config InlineBackend.figure_format = 'retina'

import matplotlib.pyplot as plt
import numpy as np
import time

import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision import datasets, transforms
import torchvision.models as models
from collections import OrderedDict

data_dir = 'flowers'
train_dir = data_dir + '/train'
valid_dir = data_dir + '/valid'
test_dir = data_dir + '/test'


train_transforms = transforms.Compose([transforms.Resize(224),
                                       transforms.RandomRotation(30),
                                       transforms.RandomResizedCrop(224),
                                       transforms.RandomHorizontalFlip(),
                                       transforms.ToTensor(),
                                       transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                                                            std=[0.229, 0.224, 0.225])])



validn_transforms = transforms.Compose([transforms.Resize(224),
                                        transforms.CenterCrop(224),
                                        transforms.ToTensor(),
                                        transforms.Normalize((0.485, 0.456, 0.406), 
                                                            (0.229, 0.224, 0.225))])

test_transforms = transforms.Compose([ transforms.Resize(224),
                                       transforms.RandomResizedCrop(224),
                                       transforms.ToTensor(),
                                       transforms.Normalize((0.485, 0.456, 0.406), 
                                                            (0.229, 0.224, 0.225))])


train_data = datasets.ImageFolder(train_dir,
                                transform=train_transforms)

validn_data = datasets.ImageFolder(valid_dir,
                                transform=validn_transforms)

test_data = datasets.ImageFolder(test_dir,
                                transform=test_transforms)



trainloader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True)
validnloader = torch.utils.data.DataLoader(validn_data, batch_size=32, shuffle=True)
testloader = torch.utils.data.DataLoader(test_data, batch_size=32, shuffle=True)


model = models.vgg16(pretrained=True)
model


for param in model.parameters():
    param.requires_grad = False

classifier = nn.Sequential(OrderedDict([            
                          ('fc1', nn.Linear(3*224*224, 10000)), 
                          ('relu', nn.ReLU()),
                          ('fc2', nn.Linear(10000, 5000)),
                          ('relu', nn.ReLU()),
                          ('fc3', nn.Linear(5000, 102)),
                          ('output', nn.LogSoftmax(dim=1))
                          ]))

model.classifier = classifier

classifier


criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=0.001)
model.cuda()

epochs = 1
steps = 0
training_loss = 0
print_every = 300
for e in range(epochs):
    model.train()
    for images, labels in iter(trainloader):
        steps == 1

        images.resize_(32,3*224*224)

        inputs = Variable(images.cuda())
        targets = Variable(labels.cuda())
        optimizer.zero_grad()

        output = model.forward(inputs)
        loss = criterion(output, targets)
        loss.backward()
        optimizer.step()

        training_loss += loss.data[0]

        if steps % print_every == 0:
            print("Epoch: {}/{}... ".format(e+1, epochs),
                  "Loss: {:.4f}".format(training_loss/print_every))

            running_loss = 0



  Traceback


ValueError                                Traceback (most recent call last)
&lt;ipython-input-17-30552f4b46e8&gt; in &lt;module&gt;()
     15         optimizer.zero_grad()
     16 
---&gt; 17         output = model.forward(inputs)
     18         loss = criterion(output, targets)
     19         loss.backward()

/opt/conda/lib/python3.6/site-packages/torchvision-0.2.0-py3.6.egg/torchvision/models/vgg.py in forward(self, x)
     39 
     40     def forward(self, x):
---&gt; 41         x = self.features(x)
     42         x = x.view(x.size(0), -1)
     43         x = self.classifier(x)

/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    323         for hook in self._forward_pre_hooks.values():
    324             hook(self, input)
--&gt; 325         result = self.forward(*input, **kwargs)
    326         for hook in self._forward_hooks.values():
    327             hook_result = hook(self, input, result)

/opt/conda/lib/python3.6/site-packages/torch/nn/modules/container.py in forward(self, input)
     65     def forward(self, input):
     66         for module in self._modules.values():
---&gt; 67             input = module(input)
     68         return input
     69 

/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    323         for hook in self._forward_pre_hooks.values():
    324             hook(self, input)
--&gt; 325         result = self.forward(*input, **kwargs)
    326         for hook in self._forward_hooks.values():
    327             hook_result = hook(self, input, result)

/opt/conda/lib/python3.6/site-packages/torch/nn/modules/conv.py in forward(self, input)
    275     def forward(self, input):
    276         return F.conv2d(input, self.weight, self.bias, self.stride,
--&gt; 277                         self.padding, self.dilation, self.groups)
    278 
    279 

/opt/conda/lib/python3.6/site-packages/torch/nn/functional.py in conv2d(input, weight, bias, stride, padding, dilation, groups)
     83     """
     84     if input is not None and input.dim() != 4:
---&gt; 85         raise ValueError("Expected 4D tensor as input, got {}D tensor instead.".format(input.dim()))
     86 
     87     f = _ConvNd(_pair(stride), _pair(padding), _pair(dilation), False,

ValueError: Expected 4D tensor as input, got 2D tensor instead.


Could it be because I am using Linear operation in my layer definitions? 

There are two problems with your network -  


You created your own classifier whose first layer accepts input of size (3*224*224), but this is not the output size of the features part of vgg16. Features output a tensor of size (25088)
  
You are resizing your input to be a tensor of shape (3*224*224) (for each batch) but the features part of vgg16 expects an input of (3, 224, 224). Your custom classifier comes after the features, so you need to prepare your input for features not for classifier.


Solution

To solve the first you need to change your definition of classifier to -  

classifier = nn.Sequential(OrderedDict([            
                          ('fc1', nn.Linear(25088, 10000)), 
                          ('relu', nn.ReLU()),
                          ('fc2', nn.Linear(10000, 5000)),
                          ('relu', nn.ReLU()),
                          ('fc3', nn.Linear(5000, 102)),
                          ('output', nn.LogSoftmax(dim=1))
                          ]))


To solve the second problem, change images.resize_(32,3*224*224) to images.resize_(32, 3, 224, 224).

P.S. - A word of advice - Your classifier's first layer output of 10000 units is very large. You should try to keep it around 4000 as done in the original classifier (Even better if you use the original weights for the first layer only, as those have proven to be good features as well over the time)
