[FIXED] RuntimeError: mat1 and mat2 shapes cannot be multiplied (1280x5 and 6400x4096)?

Issue

Defining Alexnet using the following code,I can train successfully.But when I want to see the output of each layer，it will be an error ‘RuntimeError: mat1 and mat2 shapes cannot be multiplied (1280x5 and 6400x4096)?’

class AlexNet(nn.Module):
    def __init__(self):
      super(AlexNet, self).__init__()
      self.conv = nn.Sequential(
          nn.Conv2d(1, 96, 11, 4),
          nn.ReLU(),
          nn.MaxPool2d(3, 2),
          nn.Conv2d(96, 256, 5, 1, 2),
          nn.ReLU(),
          nn.MaxPool2d(3, 2),
          nn.Conv2d(256, 384, 3, 1, 1),
          nn.ReLU(),
          nn.Conv2d(384, 384, 3, 1, 1),
          nn.ReLU(),
          nn.Conv2d(384, 256, 3, 1, 1),
          nn.ReLU(),
          nn.MaxPool2d(3, 2)
         )
      self.fc = nn.Sequential(
          nn.Linear(256*5*5, 4096),
          nn.ReLU(),
          nn.Dropout(0.5),
          nn.Linear(4096, 4096),
          nn.ReLU(),
          nn.Dropout(0.5),
          nn.Linear(4096, 10)
        )
      
    def forward(self, img):
      feature = self.conv(img)
      output = self.fc(feature.view(img.shape[0], -1))
      return output

X=torch.randn(1,1,224,224)
for name,layer in net.named_children():
  X=layer(X)
  print(name,X.shape)

Could u help me？

Solution

You forgot to flatten the output array of self.conv in the for cycle. You can split it into two cycles, one for the convolution layers, and one for the fully connected ones.

X = torch.randn(1, 1, 224, 224)
for name, layer in net.conv.named_children():
  X = layer(X)
  print(name, X.shape)

X = X.flatten()  # or X = X.view(X.shape[0], -1)

for name, layer in net.fc.named_children():
  X = layer(X)
  print(name, X.shape)

Answered By - aretor