If you want Keras style model.summary() then torchsummary is there. But it only tells you how tensors flows through your model. It does not tell you the real structure of your model (if you know what I mean).
In that case print(model) does a decent job. But it does not prints more information about model. That is where torchdescribe comes in.
-
Using
pip,pip install torchdescribe -
From command line run,
$ git clone https://github.com/vidit1999/pytorch-describe $ cd pytorch-describe $ python setup.py install
from torchdescribe import describe
input_shape = ... # input_shape(s) of model
class Model(nn.Module):
def __init__(self,...):
super(Model, self).__init__()
...
# model attribute definations
...
# model methods
...
model = Model()
describe(model, input_shape)-
import torch import torch.nn as nn from torchdescribe import describe device = 'cuda' if torch.cuda.is_available() else 'cpu' class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() self.conv1 = nn.Conv2d(1, 32, kernel_size=5) self.conv2 = nn.Conv2d(32, 32, kernel_size=5) self.conv3 = nn.Conv2d(32,64, kernel_size=5) self.fc1 = nn.Linear(3*3*64, 256) self.fc2 = nn.Linear(256, 10) def forward(self, x): x = torch.relu(self.conv1(x)) x = torch.relu(torch.max_pool2d(self.conv2(x), 2)) x = torch.relu(torch.max_pool2d(self.conv3(x),2)) x = x.view(-1,3*3*64 ) x = torch.relu(self.fc1(x)) x = self.fc2(x) return torch.log_softmax(x, dim=1) cnn = CNN().to(device) describe(model=cnn, input_shape=(1, 28, 28), batch_size=1000)
------------------------------------------------------------- CNN ------------------------------------------------------------- ============================================================= CNN( (conv1): Conv2d(1, 32, kernel_size=(5, 5), stride=(1, 1)) (conv2): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1)) (conv3): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1)) (fc1): Linear(in_features=576, out_features=256, bias=True) (fc2): Linear(in_features=256, out_features=10, bias=True) ) ============================================================= ------------------------------------------------------------- Total parameters : 228,010 Trainable parameters : 228,010 Non-trainable parameters : 0 ------------------------------------------------------------- Model device : CPU Batch size : 1000 Input shape : (1000, 1, 28, 28) Output shape : (1000, 10) Input size (MB) : 2.99 Forward/backward pass size (MB) : 257.97 Params size (MB) : 0.87 Estimated Total Size (MB) : 261.83 ------------------------------------------------------------- -
import torch from torchvision import models from torchdescribe import describe device = 'cuda' if torch.cuda.is_available() else 'cpu' vgg16 = models.vgg16().to(device) describe(vgg16, (3, 224, 224), 100)
------------------------------------------------------------------------------------ VGG ------------------------------------------------------------------------------------ ==================================================================================== 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=1000, bias=True) ) ) ==================================================================================== ------------------------------------------------------------------------------------ Total parameters : 138,357,544 Trainable parameters : 138,357,544 Non-trainable parameters : 0 ------------------------------------------------------------------------------------ Model device : CPU Batch size : 100 Input shape : (100, 3, 224, 224) Output shape : (100, 1000) Input size (MB) : 57.42 Forward/backward pass size (MB) : 21878.87 Params size (MB) : 527.79 Estimated Total Size (MB) : 22464.08 ------------------------------------------------------------------------------------ -
import torch import torch.nn as nn from torchdescribe import describe device = 'cuda' if torch.cuda.is_available() else 'cpu' class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.layers = nn.Sequential( nn.Conv2d(1, 16, kernel_size=4, stride=2, padding=1), nn.ReLU(), nn.Flatten(), ) def forward(self, x, y): a = self.layers(x) b = self.layers(y) return a, b net = Net().to(device) describe(net, [(1, 16, 16), (1, 28, 28)], 500)
------------------------------------------------------------------------- Net ------------------------------------------------------------------------- ========================================================================= Net( (layers): Sequential( (0): Conv2d(1, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) (1): ReLU() (2): Flatten(start_dim=1, end_dim=-1) ) ) ========================================================================= ------------------------------------------------------------------------- Total parameters : 272 Trainable parameters : 272 Non-trainable parameters : 0 ------------------------------------------------------------------------- Model device : CPU Batch size : 500 Input shape : [(500, 1, 16, 16), (500, 1, 28, 28)] Output shape : [(500, 1024), (500, 3136)] Input size (MB) : 1.98 Forward/backward pass size (MB) : 63.48 Params size (MB) : 0.00 Estimated Total Size (MB) : 65.46 ------------------------------------------------------------------------- -
import torch import torch.nn as nn from torchdescribe import describe device = 'cuda' if torch.cuda.is_available() else 'cpu' class RNN(nn.Module): def __init__(self): super(RNN, self).__init__() self.hidden_size = 20 self.embedding = nn.Embedding(10, 20) self.gru = nn.GRU(20, 20) def forward(self, input, hidden): embedded = self.embedding(input).view(1, 1, -1) output = embedded output, hidden = self.gru(output, hidden) return output, hidden rnn = RNN().to(device) describe(rnn, (3, 6), suppress_error=True)
-------------------------------------------------- RNN -------------------------------------------------- ================================================== RNN( (embedding): Embedding(10, 20) (gru): GRU(20, 20) ) ================================================== -------------------------------------------------- Total parameters : 2,720 Trainable parameters : 2,720 Non-trainable parameters : 0 -------------------------------------------------- Model device : CPU Batch size : 1 Input shape : (1, 3, 6) Output shape : [] Input size (MB) : -1 Forward/backward pass size (MB) : -1 Params size (MB) : 0.01 Estimated Total Size (MB) : -1 --------------------------------------------------
- For Model Size Estimation help is taken from here.
- This project is inspired by torchsummary.