Prithwish Jana
Prithwish Jana

Reputation: 317

Implementation of VGG16 on Pytorch giving size mismatch error

Snippet of my code implementation on PyTorch is:

model = models.vgg16(pretrained = False)
classifier = nn.Sequential(
        nn.Linear(25088, 128),
        nn.ReLU(True),
        nn.Dropout(),
        nn.Linear(128, 128),
        nn.ReLU(True),
        nn.Dropout(),
        nn.Linear(128, 20)
)
model.classifier = classifier

I'm feeding images of input size (60x60x3) and batch_size = 30.

When I run the code from Linux (Ubuntu) Terminal (with PyTorch Version: 1.0.0, Torchvision Version: 0.2.1) it gives me, the following error message:

RuntimeError: size mismatch, m1: [30 x 512], m2: [25088 x 128]

While, when I run it from Spyder (Anaconda) on Windows (with PyTorch Version: 1.0.1, Torchvision Version: 0.2.2), it runs perfectly.

Am I missing something or is this because of some version mismatch in Pytorch and Torchvision? Both, I'm running on Python 3.6. Please suggest.

[UPDATE: Mistakenly interchanged the version numbers for the error-case and error-free case. Thanks @Manoj Mohan for pointing it out]

Upvotes: 2

Views: 863

Answers (1)

Manoj Mohan
Manoj Mohan

Reputation: 6034

It's probably the other way around. Things run perfectly on torchvision 0.2.2 and fails on torch vision 0.2.1.

This change of using AdaptiveAvgPool2d that went into 0.2.2 is why you don't see the error. https://github.com/pytorch/vision/commit/83b2dfb2ebcd1b0694d46e3006ca96183c303706

>>> import torch
>>> model = models.vgg16(pretrained = False)
>>> x = torch.randn(1,3,60,60) # random image
>>> feat = model.features(x)
>>> flat_feat = feat.view(feat.size(0), -1) # flatten
>>> flat_feat.shape
torch.Size([1, 512])
>>> model.classifier(flat_feat)

RuntimeError: size mismatch, m1: [1 x 512], m2: [25088 x 4096] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:940

You see the error of size mismatch. After, adaptive average pooling, things work fine.

>>> import torch.nn.functional as F
>>> avg = F.adaptive_avg_pool2d(feat, (7,7))
>>> avg = avg.view(avg.size(0), -1)
>>> output = model.classifier(avg)
>>> output.shape
torch.Size([1, 1000])

Upvotes: 1

Related Questions