I’m trying to build a CNN for classifying medical images into 5 classes. Here’s the architecture i used:
class ConvNet(nn.Module):
def __init__(self):
super(ConvNet, self).__init__()
self.conv1 = nn.Conv2d(3, 32, 3)
self.conv2 = nn.Conv2d(32, 32, 3)
self.pool1 = nn.MaxPool2d(2, 2)
self.conv3 = nn.Conv2d(32, 64, 3)
self.conv4 = nn.Conv2d(64, 64, 3)
self.pool2 = nn.MaxPool2d(2, 2)
self.conv5 = nn.Conv2d(64, 128, 3)
self.conv6 = nn.Conv2d(128, 128, 3)
self.pool3 = nn.MaxPool2d(2, 2)
self.conv7 = nn.Conv2d(128, 256, 3)
self.conv8 = nn.Conv2d(256, 256, 3)
self.pool4 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(256*12*12,4096)
self.fc2 = nn.Linear(4096, 5)
def forward(self, input):
x = self.pool1(self.conv2(F.relu(self.conv1(input))))
x = self.pool2(self.conv4(F.relu(self.conv3(x))))
x = self.pool3(self.conv6(F.relu(self.conv5(x))))
x = self.pool4(self.conv8(F.relu(self.conv7(x))))
#flatten the tensor for the FC
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
the input from trainloader is (10,3,256,256) and labels are of size “torch.Size([10])”.
and here is a snippet from the main:
for idx, batch in enumerate(trainloader):
data = batch['image']
labels = batch['label']
if use_cuda:
data = data.cuda()
labels = labels.cuda()
if labels.data.size()[0] == batch_size:
optimizer.zero_grad()
output = model(data)
loss = criterion(output, labels)
loss.backward()
optimizer.step()
this line loss = criterion(output, labels.data) gives me error that error because they are not the same dimensions: input and target shapes do not match: input [10 x 5], target [10 x 1]
according to my understanding to cnn and the softmax layer at the end, the output should be 10x1(which is the highest class’s probability among those five ) but i found it to be 10x5. how can i fix this?
Thanks a lot