AlexNet: Output shape mismatch between Conv2D and dense layer

minimafinder: <p>I referenced <a href="https://proceedings.neurips.cc/paper_files/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf" rel="noopener nofollow ugc">Krizhevsky et al. (2012)</a> and attempted to replicate the model as defined in <em>Figure 2</em>. To test the model, I am passing a subset of a small number of images as tensors one at a time. When I run the model, I get the following error:</p>
<p><code>RuntimeError: linear(): input and weight.T shapes cannot be multiplied (256x10 and 9216x2048)</code></p>
<p>This is happening because the outputs from the fifth convolutional layer don’t have the right dimensions needed for the subsequent dense layer. <strong>Is there a way to remedy this situation without toying too much with the parameter sizes? What am I doing wrong here?</strong></p>
<p>(While creating the model, I referenced <a href="https://www.cs.toronto.edu/~rgrosse/courses/csc321_2018/tutorials/tut6_slides.pdf" rel="noopener nofollow ugc">slide 11 from this lecture deck</a> so that the model size matches the size of the original model. I am trying to recreate what the authors did in the paper without ad-hoc modifications; they unfortunately do not mention all the granular model details, though.)</p>
<h3><a name="relevant-code-below-1" class="anchor" href="#relevant-code-below-1"></a>Relevant code below</h3>
<pre><code class="lang-auto">class AlexNet(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1_top = nn.Conv2d(
            in_channels=3,
            out_channels=48,
            kernel_size=11,
            stride=4,
        ).to(device1)

        self.conv1_bottom = nn.Conv2d(
            in_channels=3,
            out_channels=48,
            kernel_size=11,
            stride=4,
        ).to(device2)

        self.conv2_top = nn.Conv2d(
            in_channels=48,
            out_channels=128,
            kernel_size=5,
            padding=2,
        ).to(device1)

        self.conv2_bottom = nn.Conv2d(
            in_channels=48,
            out_channels=128,
            kernel_size=5,
            padding=2,
        ).to(device2)

        self.conv3_top = nn.Conv2d(
            in_channels=256,
            out_channels=192,
            kernel_size=3,
            padding=1, 
        ).to(device1)

        self.conv3_bottom = nn.Conv2d(
            in_channels=256,
            out_channels=192,
            kernel_size=3,
            padding=1,
        ).to(device2)

        self.conv4_top = nn.Conv2d(
            in_channels=192,
            out_channels=192,
            kernel_size=3,
            padding=1,
        ).to(device1)

        self.conv4_bottom = nn.Conv2d(
            in_channels=192,
            out_channels=192,
            kernel_size=3,
            padding=1,
        ).to(device2)

        self.conv5_top = nn.Conv2d(
            in_channels=192,
            out_channels=128,
            kernel_size=3,
            padding=1,
        ).to(device1)

        self.conv5_bottom = nn.Conv2d(
            in_channels=192,
            out_channels=128,
            kernel_size=3,
            padding=1,
        ).to(device2)

        self.dense6_top = nn.Linear(
            in_features=256 * 6 * 6, 
            out_features=2048,
        ).to(device1)

        self.dense6_bottom = nn.Linear(
            in_features=256 * 6 * 6,
            out_features=2048,
        ).to(device2)

        self.dense7_top = nn.Linear(
            in_features=4096, 
            out_features=2048,
        ).to(device1)

        self.dense7_bottom = nn.Linear(
            in_features=4096, 
            out_features=2048,
        ).to(device2)

        self.dense_last = nn.Linear(
            in_features=4096, 
            out_features=1000,
        ).to(device1)

        # local response normalization layer with 
        # hyperparameters as described in the paper
        self.lrn = nn.LocalResponseNorm(
            size=5,
            alpha=0.0001, 
            beta=0.75, 
            k=2
        )

        # not all layer details are mentioned in the paper
        # For the kernel size, I downloaded and referenced 
        # the model wiki provided by the authors
        # see file "LayerParams.wiki.ini"
        self.maxpool = nn.MaxPool2d(
            kernel_size=3,
            stride=2,
        )

        self.dropout = nn.Dropout(
            p=0.5,
        )

    
    def forward(self, x):
        top_image = x[:, :int(x.shape[1] / 2), :].to(device1)
        bottom_image = x[:, int(x.shape[1] / 2):, :].to(device2)

        # first layer
        top_image = nnF.relu(self.conv1_top(top_image))
        bottom_image = nnF.relu(self.conv1_bottom(bottom_image))

        ## response norm after first layer
        top_image = self.lrn(top_image)
        bottom_image = self.lrn(bottom_image)

        ## max pooling after first response norm
        top_image = self.maxpool(top_image)
        bottom_image = self.maxpool(bottom_image)

        
        # second layer
        top_image = nnF.relu(self.conv2_top(top_image))
        bottom_image = nnF.relu(self.conv2_bottom(bottom_image))

        ## response norm after second layer
        top_image = self.lrn(top_image)
        bottom_image = self.lrn(bottom_image)

        ## max pooling after second response norm
        top_image = self.maxpool(top_image)
        bottom_image = self.maxpool(bottom_image)

        # my best interpretation of how the output is passed to the third layer
        image_full = torch.cat(
            (top_image.to(device2), bottom_image.to(device2)), 
            dim=0
        )

        print(top_image.shape)
        print(f"full before layer 3:", image_full.shape)

        # third layer
        top_image = nnF.relu(self.conv3_top(image_full.to(device1)))
        bottom_image = nnF.relu(self.conv3_bottom(image_full.to(device2)))


        # fourth layer
        top_image = nnF.relu(self.conv4_top(top_image))
        bottom_image = nnF.relu(self.conv4_bottom(bottom_image))

        
        # fifth layer
        top_image = nnF.relu(self.conv5_top(top_image))
        bottom_image = nnF.relu(self.conv5_bottom(bottom_image))

        ## max pooling after the fifth layer
        print("fifth layer shape before max pooling:", top_image.shape)
        print("fifth layer shape before max pooling:", bottom_image.shape)
        
        # my best interpretation of how the output is passed to the sixth layer
        image_full = torch.cat(
            (top_image.to(device1), bottom_image.to(device1)), 
            dim=0
        )
        
        image_full = self.maxpool(image_full)
        #image_full = self.averagepool(image_full)
        print("fifth layer shape after max pooling:", image_full.shape)
        
        # reshaping the input for the dense layer
        image_full = image_full.reshape(image_full.size(0), -1)
        
        # sixth layer
        image_full = self.dropout(image_full)
        top_image = nnF.relu(self.dense6_top(image_full.to(device1)))
        #bottom_image = self.dropout(bottom_image)
        bottom_image = nnF.relu(self.dense6_bottom(image_full.to(device2)))


        image_full = torch.cat(
            (top_image.to(device2), bottom_image.to(device2)), 
            dim=0
        )

        print(f"full before layer 7:", image_full.shape)


        # seventh layer
        image_full = self.dropout(image_full)
        top_image = nnF.relu(self.dense7_top(image_full.to(device1)))
        #bottom_image = self.dropout(bottom_image)
        bottom_image = nnF.relu(self.dense7_bottom(image_full.to(device2)))

        
        # final layer
        combined_output = torch.cat(
            (top_image.to(device1), bottom_image.to(device1)), 
            dim=0
        )
        
        x = F.softmax(self.dense_last(combined_output))
        
        return x

alexnet = AlexNet()
print(alexnet)
</code></pre>
<p>And then I run it as:</p>
<pre><code class="lang-auto">num_classes = 1000
learning_rate = 0.005
momentum=0.9
weight_decay=0.005

# defining the loss and the optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(
    alexnet.parameters(), 
    lr=learning_rate,
    momentum=momentum,
    weight_decay=weight_decay,
)

steps_to_take = len(train_subset_loader)

num_epochs = 10

for epoch in range(num_epochs):
    for i, (inputs, labels) in enumerate(train_subset_loader):  
        # forward pass
        print(inputs.shape)
        outputs = alexnet(inputs)
        loss = criterion(outputs, labels)
        
        # backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(
        f"Epoch [{epoch + 1}/{num_epochs}], "
        f"Step [{i + 1}/{steps_to_take}], " 
        f"Loss: {loss.item():.4f}"
    )
</code></pre>
<p><strong>The output</strong></p>
<pre><code class="lang-auto">torch.Size([3, 224, 224])
torch.Size([128, 5, 12])
full before layer 3: torch.Size([256, 5, 12])
fifth layer shape before max pooling: torch.Size([128, 5, 12])
fifth layer shape before max pooling: torch.Size([128, 5, 12])
fifth layer shape after max and avg pooling: torch.Size([256, 2, 5])
full before layer 6: torch.Size([256, 10])

(a bunch of error traceback)

RuntimeError: linear(): input and weight.T shapes cannot be multiplied (256x10 and 9216x2048)
</code></pre>

ptrblck: <p>My follow up post shows the plain model implementation without the adaptive pooling layer, which works fine for me: <a href="https://discuss.pytorch.org/t/why-does-alexnet-in-torch-vision-use-average-pooling/84313/7" class="inline-onebox">Why does Alexnet in torch vision use Average Pooling - #7 by ptrblck</a></p>

minimafinder: <p>Thank you; I saw the reply, and it was helpful quick check. A key difference, though, is that I am splitting the model across two GPUs <strong>and I was previously also splitting the input image vertically</strong>, which changed the output shapes through the model.</p>
<p>(I was splitting the input because of a pretty literal interpretation of the following statement in the paper: “One GPU runs the layer-parts at the <em>top of the figure</em> while the other runs the layer-parts at the bottom.”)</p>
<p>If the image is vertically split (which I think is incorrect), then Adaptive average pooling is the only way to make the model work. Now, instead splitting the input, I pass the entire image to both the top and bottom layers, <strong>but for the top layer, the bottom half of the figure is masked to zero and vice-versa.</strong> The outputs of layer 3, 6, and 7 are concatenated before being passed to the subsequent layer. With this revised input scheme, AlexNet with maxpool works great!</p>
<p>I am still unsure how exactly the authors’ passed the input in the original paper…</p>