RuntimeError: mat1 and mat2 shapes cannot be multiplied (5x46656 and 50176x3)

user_18168227 0
Viewed 105

I am trying to build a cnn model and while building i am getting the run time error. I have image size of 224x224x3 and batch size of 5 and i have 3 classes to be predicted.

class CNN(nn.Module):
def __init__(self):
    super().__init__()
    #5*3*224*224
    self.conv1=nn.Conv2d(3,6,3,1)
    self.relu1=nn.ReLU()
    self.pool=nn.MaxPool2d(2)
    #112*112*6

    self.conv2=nn.Conv2d(6,16,3,1)
    self.relu2=nn.ReLU()

    self.fc1=nn.Linear(16*56*56,3)

def forward(self,x):
    x=self.pool(self.relu1(self.conv1(x)))
    x=self.pool(self.relu2(self.conv2(x)))
    x=x.flatten(1)
    x=self.fc1(x)
    return x

I get the run time error during training. How do I fix this?

conv-neural-network pytorch python
1 Answers

Considering the size of your input, your fully connected layer should have 16*54*54 neurons, not 16*56*56.

self.fc1 = nn.Linear(16*54*54, 3)

Alternatively, you can you a lazy module that will infer the number of neurons needed: nn.LazyLinear:

self.fc1 = nn.LazyLinear(3)
user_63313690
Related