sksq96 / pytorch-summary

Model summary in PyTorch similar to `model.summary()` in Keras
MIT License
4.02k stars 413 forks source link

Wrong result on Dropout layer #81

Open kumarneelabh13 opened 5 years ago

kumarneelabh13 commented 5 years ago
import torchvision.models as models
from torchsummary import summary

mobilenet = models.mobilenet_v2()
summary(mobilenet, (3, 224, 224))

On running this code, I got an error in the output shape of the Dropout layer. Here are the last few lines from the output :

InvertedResidual-153 [-1, 320, 7, 7] 0 Conv2d-154 [-1, 1280, 7, 7] 409,600 BatchNorm2d-155 [-1, 1280, 7, 7] 2,560 ReLU6-156 [-1, 1280, 7, 7] 0 Dropout-157 [-1, 1280] 0 Linear-158 [-1, 1000] 1,281,000

Total params: 3,504,872 Trainable params: 3,504,872 Non-trainable params: 0

The result for the Dropout layer is wrong because Dropout does not change the shape of the input as per the PyTorch docs.

frgfm commented 4 years ago

This is not an error from torchsummary, rather a lack of information on the functional API of torch from torchsummary.

You pointed out that Dropout does not change the shape of the input, which is correct. But you assume that the input to the Dropout layer is the output of ReLU6-156. Always check the ̀forward` implementation of the model ;)

In this case, there is an adaptive average pooling and a .reshape(x.shape[0], -1) over here: https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenet.py#L155

There are not picked up by the hooking mechanism, hence your output!