How can I check the image sizes while CNN?

Question

I'm trying to classify cat and dog in CNN with PyTorch. While I made few layers and processing images, I found that final processed feature map size doesn't match with calculated size. So I tried to check feature map size step by step in CNN process with print shape but it doesn't work. I heard tensorflow enables check tensor size in steps but how can I do that?

What I want is:

        def __init__(self):
        super(CNN, self).__init__()
        conv1 = nn.Conv2d(1, 16, 3, 1, 1)
        conv1_1 = nn.Conv2d(16, 16, 3, 1, 1)
        pool1 = nn.MaxPool2d(2)
        conv2 = nn.Conv2d(16, 32, 3, 1, 1)
        conv2_1 = nn.Conv2d(32, 32, 3, 1, 1)
        pool2 = nn.MaxPool2d(2)
        conv3 = nn.Conv2d(32, 64, 3, 1, 1)
        conv3_1 = nn.Conv2d(64, 64, 3, 1, 1)
        conv3_2 = nn.Conv2d(64, 64, 3, 1, 1)
        pool3 = nn.MaxPool2d(2)

        self.conv_module = nn.Sequential(
            conv1,
            nn.ReLU(),
            conv1_1,
            nn.ReLU(),
            pool1,
            # check first result size
            conv2,
            nn.ReLU(),
            conv2_1,
            nn.ReLU(),
            pool2,
            # check second result size
            conv3,
            nn.ReLU(),
            conv3_1,
            nn.ReLU(),
            conv3_2,
            nn.ReLU(),
            pool3,
            # check third result size
            pool4,
            # check fourth result size
            pool5
            # check fifth result size
        )

If there's any other way to check feature size at every step, please give some advice. Thanks in advance.

Answer 1

To do that you shouldn't use nn.Sequential . Just initialize your layers in __init__() and call them in the forward function. In the forward function you can print the shapes out. For example like this:

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        
        self.conv1 = nn.Conv2d(...)
        self.maxpool1 = nn.MaxPool2d()

        self.conv2 = nn.Conv2d(...)
        self.maxpool2 = nn.MaxPool2d()
    
    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.maxpool1(x)
        print(x.size())

        x = self.conv2(x)
        x = F.relu(x)
        x = self.maxpool2(x)
        print(x.size())

Hope thats what you looking for!