I m trying to write a simple custom neural network for the image classification task. While running model I am getting the runtime error. I changed linear layer input either 21168/28224 but the issue still exists.
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
# first layer
self.conv1 = nn.Conv2d(3, 20, 1)
nn.init.kaiming_uniform_(self.conv1.weight, nonlinearity='relu')
self.act1 = nn.ReLU()
self.pool1 = nn.MaxPool2d((2,2), stride=(2,2))
# second layer
self.conv2 = nn.Conv2d(3, 40, 1)
nn.init.kaiming_uniform_(self.conv2.weight, nonlinearity='relu')
self.act2 = nn.ReLU()
self.pool2 = nn.MaxPool2d((2,2), stride=(2,2))
# third layer
self.conv3 = nn.Conv2d(3, 40, 3)
nn.init.kaiming_uniform_(self.conv2.weight, nonlinearity='relu')
self.act3 = nn.ReLU()
self.pool3 = nn.MaxPool2d((2,2), stride=(2,2))
# fourth layer
self.conv4 = nn.Conv2d(3, 40, 5)
nn.init.kaiming_uniform_(self.conv2.weight, nonlinearity='relu')
self.act4 = nn.ReLU()
self.pool4 = nn.MaxPool2d((2,2), stride=(2,2))
# self.fc1 = nn.Linear(20*110*110, 3)
# self.fc1 = nn.Linear(141120, 3)
self.fc1 = nn.Linear(28224, 3)
def forward(self, x):
x1 = self.conv1(x)
x1 = self.act1(x)
x1 = self.pool1(x)
x2 = self.conv2(x1)
x2 = self.act2(x1)
x2 = self.pool2(x1)
x3 = self.conv3(x1)
x3 = self.act3(x1)
x3 = self.pool3(x1)
x4 = self.conv4(x1)
x4 = self.act4(x1)
x4 = self.pool4(x1)
x = torch.cat([x2.view(-1, 4), x3.view(-1, 4), x4.view(-1, 4)], -2)
x = x.view(4, -1)
x = F.relu(self.fc1(x))
return F.log_softmax(x)
criterion = nn.CrossEntropyLoss()
# Observe that all parameters are being optimized
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
model = Model()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model_tr = train_model(model, criterion, optimizer_ft, exp_lr_scheduler, 4)
Following is the stack trace:
RuntimeError Traceback (most recent call last)
<ipython-input-38-8d53b5624a06> in <module>()
----> 1 model_tr = train_model(model, criterion, optimizer_ft, exp_lr_scheduler, 4)
5 frames
<ipython-input-14-1a5d4f162548> in train_model(model, criterion, optimizer, scheduler, num_epochs)
30 # track history if only in train
31 with torch.set_grad_enabled(phase == 'train'):
---> 32 outputs = model(inputs)
33 _, preds = torch.max(outputs, 1)
34 loss = criterion(outputs, labels)
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
530 result = self._slow_forward(*input, **kwargs)
531 else:
--> 532 result = self.forward(*input, **kwargs)
533 for hook in self._forward_hooks.values():
534 hook_result = hook(self, input, result)
<ipython-input-36-2714cf86e78a> in forward(self, x)
56
57 x = x.view(4, -1)
---> 58 x = F.relu(self.fc1(x))
59 return F.log_softmax(x)
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
530 result = self._slow_forward(*input, **kwargs)
531 else:
--> 532 result = self.forward(*input, **kwargs)
533 for hook in self._forward_hooks.values():
534 hook_result = hook(self, input, result)
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/linear.py in forward(self, input)
85
86 def forward(self, input):
---> 87 return F.linear(input, self.weight, self.bias)
88
89 def extra_repr(self):
/usr/local/lib/python3.6/dist-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: [4 x 21168], m2: [28224 x 3] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:290
Can someone please help ?