# VGG16 Model Outputs Incorrect dimension - Transfer Learning

**URL:** <https://discuss.pytorch.org/t/vgg16-model-outputs-incorrect-dimension-transfer-learning/109805>\
**Category:** vision\
**Created:** [January 23, 2021, 9:33pm UTC](https://discuss.pytorch.org/t/vgg16-model-outputs-incorrect-dimension-transfer-learning/109805 "2021-01-23T21:33:06Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Wanderer3](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/w/dc4da7/32.png) [@Wanderer3](https://discuss.pytorch.org/u/Wanderer3)\
**Post date:** [January 23, 2021, 9:33pm UTC](https://discuss.pytorch.org/t/vgg16-model-outputs-incorrect-dimension-transfer-learning/109805/1 "2021-01-23T21:33:06Z")

</div>

I am trying to use a pre-trained VGG16 model to classify CIFAR10 on pyTorch. The model was originally trained on ImageNet.

Here is how I imported and modified the model:

```
from torchvision import models

model = models.vgg16(pretrained=True).cuda()
model.classifier[6].out_features = 10 

```

and this is the summary of the model

```
print(model)

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=10, bias=True)
  )
)

```

Now I wanted to train the model on CIFAR10, I created the train loader with batch size = 128.  
Below is the training function and I added some printing statements to check if everything is working fine or not. **The problem is the model outputs for each data point a 1000 prediction (as the original - unmodified version).**

```
def train(model, optimizer, train_loader, epoch=5):
    """
    This function updates/trains client model on client data
    """
    model.train()
    for e in range(epoch):
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            print("shape output: ", output.shape)
            probs = F.softmax(output, dim=1) #get the entropy 
            print("entropy: ", probs)
            _, predicted = torch.max(output.data, 1)
            loss = criterion(output, target)
            loss.backward()
            if batch_idx % 100 == 0: # print every 100 mini-batches
                print('[%d, %5d] loss: %.3f' %
                  (e + 1, batch_idx + 1, loss.item()))
            optimizer.step()
    return loss.item()

```

Here is part of the the output which shows the output shape:

```
shape output: torch.Size([128, 1000])

```

**I don’t understand why the model outputs the results in this way. Is there something that I am missing?**

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [January 23, 2021, 9:36pm UTC](https://discuss.pytorch.org/t/vgg16-model-outputs-incorrect-dimension-transfer-learning/109805/2 "2021-01-23T21:36:48Z")

</div>

Changing the `out_features` attribute of the linear layer will not recreate the parameters, so you would need to assign a new layer instead:

```python
model = models.vgg16(pretrained=True)
in_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(in_features, out_features)

```