Hi!
I am trying to build a Convolutional Neural Network and am facing a problem. The network has 3 Conv. layers, ReLU activated, pooled, batch normalized and then flattened. Here’s the model code:
in_features = 152 #in_features for Flatten(linear) layer calculated

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 2, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(num_features=6)
        self.conv2 = nn.Conv2d(6, 8, 2, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(num_features=8)
        self.conv3 = nn.Conv2d(8, 12, 2, stride=1, padding=1)
        self.bn3 = nn.BatchNorm2d(num_features=12)

        self.fc1   = nn.Linear(in_features, 36)
        self.fc2   = nn.Linear(36, 3)
        
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.max_pool2d(out, 2)
        
        out = F.relu(self.bn2(self.conv2(out)))
        out = F.max_pool2d(out, 2)
        
        out = F.relu(self.bn3(self.conv3(out)))
        out = F.max_pool2d(out, 2)
        
        out = out.view(out.size(0), -1)
        out = F.relu(self.fc1(out))
        out = F.softmax(self.fc2(out))
        return out

pytorch_model = Net()

I am trying to do multiclass classification and so using ‘Adam’ Optimizer and ‘CrossEntropyLoss’ function.
My input is of the size: (143256,1, 150, 3) and I am trying to train using the following code:
batch_size = 100
epochs = 10

pytorch_model = pytorch_model.double()

for epoch in range(epochs):
    for i in tqdm(range(0, len(X_TRAIN), batch_size)):
        batch_x = X_TRAIN[i:i+batch_size]
        batch_y = Y_TRAIN[i:i+batch_size]
        
        optimizer.zero_grad() # clears all the gradients
        
        outputs = pytorch_model(batch_x.double())
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()
    print(loss)

And when trying to train the model I get this error:
RuntimeError                              Traceback (most recent call last)
<ipython-input-189-e1854e4cd77a> in <module>
---> 15         outputs = pytorch_model(batch_x.double())
~/anaconda3/envs/LearningML/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
--> 541             result = self.forward(*input, **kwargs)
<ipython-input-184-d3ab227f7914> in forward(self, x)
---> 27         out = F.relu(self.fc1(out))
~/anaconda3/envs/LearningML/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
--> 541             result = self.forward(*input, **kwargs)
~/anaconda3/envs/LearningML/lib/python3.6/site-packages/torch/nn/modules/linear.py in forward(self, input)
---> 87         return F.linear(input, self.weight, self.bias)
~/anaconda3/envs/LearningML/lib/python3.6/site-packages/torch/nn/functional.py in linear(input, weight, bias)
   1368     if input.dim() == 2 and bias is not None:
   1369         # fused op is marginally faster
-> 1370         ret = torch.addmm(bias, input, weight.t())
   1371     else:
   1372         output = input.matmul(weight.t())

RuntimeError: size mismatch, m1: [100 x 228], m2: [152 x 36] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:197

I am fairly new to using PyTorch, so I know there’s something I am doing wrong here. Can anyone please explain what is that? Any help would be really appreciated…
Thanks!