How do I get the summary of a PyTorch model?

Viewed 118

I am trying to load a CNN pytorch model using:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1d = nn.Conv1d(in_channels=1, out_channels=32, kernel_size=7, stride=1, padding=4)
        self.conv1d2 = nn.Conv1d(in_channels=32, out_channels=64, kernel_size=7, stride=1, padding=4)
        self.conv1d3 = nn.Conv1d(in_channels=64, out_channels=128, kernel_size=7, stride=1, padding=4)
        self.maxpool = nn.MaxPool1d(2)
        self.fc = nn.Linear(128, 16)

    def forward(self, x):
        x = x.unsqueeze(1)
        # layer 1
        x = self.conv1d(x)
        x = F.relu(x)
        x = self.maxpool(x)
        # # layer2
        x = self.conv1d2(x)
        x = F.relu(x)
        x = self.maxpool(x)
        # layer3
        x = self.conv1d3(x)
        x = F.relu(x)
        x = self.maxpool(x)
        # flatten
        x = x.squeeze()

        # feed to a fully-connected layer
        x = self.fc(x)
        x = torch.sigmoid(x)

        return x

net = CNN()

However, when I load Pytorch's summary function:

summary(net, (1,28))

I get a dimension error

RuntimeError: Expected 3-dimensional input for 3-dimensional weight [32, 1, 7], but got 4-dimensional input of size [2, 1, 1, 28] instead

Please any help would be grateful, thank you

1 Answers

i just change the max-pooling layer to 4 and x.squeeze(0)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1d = nn.Conv1d(in_channels=1, out_channels=32, kernel_size=7, stride=1, padding=4)
        self.conv1d2 = nn.Conv1d(in_channels=32, out_channels=64, kernel_size=7, stride=1, padding=4)
        self.conv1d3 = nn.Conv1d(in_channels=64, out_channels=128, kernel_size=7, stride=1, padding=4)
        self.maxpool = nn.MaxPool1d(4)
        self.fc = nn.Linear(128, 16)

    def forward(self, x):
        x = x.squeeze(0)
        # layer 1
        x = self.conv1d(x)
        x = F.relu(x)
        x = self.maxpool(x)
        # # layer2
        x = self.conv1d2(x)
        x = F.relu(x)
        x = self.maxpool(x)
        # layer3
        x = self.conv1d3(x)
        x = F.relu(x)
        x = self.maxpool(x)
        # flatten
        x = x.squeeze()
        # feed to a fully-connected layer
        x = self.fc(x)
        x = torch.sigmoid(x)

        return x

net = CNN()
summary(net, (1,28))
Related