In order to use torchsummary type:
from torchsummary import summary
Install it first if you don't have it.
pip install torchsummary
And then you can try it, but note from some reason it is not working unless I set model to cuda alexnet.cuda:
from torchsummary import summary
help(summary)
import torchvision.models as models
alexnet = models.alexnet(pretrained=False)
alexnet.cuda()
summary(alexnet, (3, 224, 224))
print(alexnet)
The summary must take the input size and batch size is set to -1 meaning any batch size we provide.
If we set summary(alexnet, (3, 224, 224), 32) this means use the bs=32.
summary(model, input_size, batch_size=-1, device='cuda')
Out:
Help on function summary in module torchsummary.torchsummary:
summary(model, input_size, batch_size=-1, device='cuda')
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [32, 64, 55, 55] 23,296
ReLU-2 [32, 64, 55, 55] 0
MaxPool2d-3 [32, 64, 27, 27] 0
Conv2d-4 [32, 192, 27, 27] 307,392
ReLU-5 [32, 192, 27, 27] 0
MaxPool2d-6 [32, 192, 13, 13] 0
Conv2d-7 [32, 384, 13, 13] 663,936
ReLU-8 [32, 384, 13, 13] 0
Conv2d-9 [32, 256, 13, 13] 884,992
ReLU-10 [32, 256, 13, 13] 0
Conv2d-11 [32, 256, 13, 13] 590,080
ReLU-12 [32, 256, 13, 13] 0
MaxPool2d-13 [32, 256, 6, 6] 0
AdaptiveAvgPool2d-14 [32, 256, 6, 6] 0
Dropout-15 [32, 9216] 0
Linear-16 [32, 4096] 37,752,832
ReLU-17 [32, 4096] 0
Dropout-18 [32, 4096] 0
Linear-19 [32, 4096] 16,781,312
ReLU-20 [32, 4096] 0
Linear-21 [32, 1000] 4,097,000
================================================================
Total params: 61,100,840
Trainable params: 61,100,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 18.38
Forward/backward pass size (MB): 268.12
Params size (MB): 233.08
Estimated Total Size (MB): 519.58
----------------------------------------------------------------
AlexNet(
(features): Sequential(
(0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))
(1): ReLU(inplace)
(2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
(3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
(4): ReLU(inplace)
(5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
(6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(7): ReLU(inplace)
(8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(9): ReLU(inplace)
(10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(inplace)
(12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(avgpool): AdaptiveAvgPool2d(output_size=(6, 6))
(classifier): Sequential(
(0): Dropout(p=0.5)
(1): Linear(in_features=9216, out_features=4096, bias=True)
(2): ReLU(inplace)
(3): Dropout(p=0.5)
(4): Linear(in_features=4096, out_features=4096, bias=True)
(5): ReLU(inplace)
(6): Linear(in_features=4096, out_features=1000, bias=True)
)
)