Pytorch linear/affine layer parameters confusing

user_12210238 0
Viewed 1085

I'm on the Pytorch documentation (https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html) and I'm not really understanding why they are making the the affine layer (16 * 6 * 6, 120). I understand that the last outputs from the convolution layer were 16 and the output here is 120, but even with their annotation, I'm not understanding where the 6 * 6 comes from. Can someone please explain?

import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 6 * 6, 120)  # 6*6 from image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # If the size is a square you can only specify a single number
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


net = Net()
print(net)
pytorch neural-network
1 Answers

The 6x6 comes from the height and width of x after it has been passed through your convolutions and maxpools.

Here is a simplified version where you can see how the shape changes at each point. It may help to print out the shapes in their example so you can see exactly how everything changes.

import torch
import torch.nn as nn
import torch.nn.functional as F

conv1 = nn.Conv2d(1, 6, 3)
conv2 = nn.Conv2d(6, 16, 3)

# Making a pretend input similar to theirs.
# We define an input with 1 batch, 1 channel, height 32, width 32
x = torch.ones((1,1,32,32))


# Simulating forward()

x = F.max_pool2d(F.relu(conv1(x)), (2, 2))
print(x.shape)  # torch.Size([1, 6, 15, 15])  1 batch, 6 channels, height 15, width 15

x = F.max_pool2d(F.relu(conv2(x)), 2)
print(x.shape)  # torch.Size([1, 16, 6, 6])  1 batch, 16 channels, height 6, width 6

Next they flatten x and pass it through fc1 which accepts 16*6*6 and produces 120 outputs.

user_3009080
Related