pytorch Crossentropy results in unmatched batch size

Question

I was loading a image folder using dataloader

The image folder consists of three categories(labels) ie

'/root/ant/dsd.png'

'/root/ant/sfds.png'

...

....

'/root/bee/dsf.png'

....

..

'/root/whey/sfd.png'

Here there are three classes ant,bee,whey

By executing the above code i getting an error of unmatched bacth size of output and target at criterion

Error:Expected input batch_size (3) to match target batch_size (1).

I thought that the error might be in trainloader because extracting labels of different unmatched shape

data_transform = transforms.Compose([
        transforms.Resize(size=28),
        transforms.ToTensor()
    ])
     kumda_dataset = datasets.ImageFolder(root='/content/gdrive/My Drive/Colab Notebooks/images',
                                           transform=data_transform)
#train & test 
train_size = int(0.8 * len(kumda_dataset))
test_size = len(kumda_dataset) - train_size

#splitting
train_dataset, test_dataset = torch.utils.data.random_split(kumda_dataset, [train_size, test_size])

trainloader = torch.utils.data.DataLoader(train_dataset , batch_size = 1, shuffle = True)
testloader = torch.utils.data.DataLoader(test_dataset , batch_size = 4, shuffle = False )

model=nn.Linear(784,1)
criterion=nn.CrossEntropyLoss()
optimizer=optim.SGD(model.parameters(),lr=0.01)

num=10
for epoch in range(num):
  for i,(images,labels) in enumerate(trainloader):
    images=images.reshape(-1,28*28)
    output=model(images)
    loss=criterion(output,labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if(i+1%70==0):
        print("Epoch: {}/{}, \tIteration: {}/{}, \tLoss: {}".format(epoch + 1, num, i + 1,len(dataset_loader), loss.item()))

Thank you in advance for solving

Answer 1

Your linear layer is only outputting a single value per batch item. CrossEntropyLoss expects one output dimension per class. Change to

model=nn.Linear(784, 3)

since you have 3 classes.