RuntimeError: Expected floating point type for target with class probabilities, got Long when using Cross Entropy Loss

vision
1

I am getting this error when using CrossEntropyLoss().
I am only training a linear layer and use a pre-trained model.
I also tried sending output through softmax but still got the same error.

cross_entropy_loss = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(linear_layer.parameters(), lr=3e-4)

    for epoch_idx in range(args.linear_max_epoch):
        linear_layer.train()
        loss_sum = 0
        total, correct = 0, 0

        for data, gt_labels in train_loader:
            data, gt_labels = data.to(device).float(), gt_labels.to(device)
            # data shape  torch.Size([512, 512])  [batch size, features]
            gt_labels = fn.one_hot(gt_labels.long(), num_classes = 5)
            # target shape  torch.Size([512, 5])
            output = linear_layer(data)   
            # tried output = linear_layer(data).softmax(dim = 1) and still got error
            # output shape  torch.Size([512, 5])
            # output tensor contains probabilities (output from a linear_layer)
            loss = cross_entropy_loss(output, gt_labels)   # THIS IS WHERE I GET THE ERROR
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            loss_sum += loss.detach().cpu().numpy()

            _, predicted = torch.max(output.data, 1)
            total += gt_labels.size(0)
            correct += (predicted == gt_labels).sum().item()

The error is,

Traceback (most recent call last):
  File "tools/linear_eval.py", line 156, in <module>
    linear_eval_train(1)
  File "tools/linear_eval.py", line 104, in linear_eval_train
    loss = cross_entropy_loss(output, gt_labels)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/loss.py", line 1166, in forward
    label_smoothing=self.label_smoothing)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py", line 3014, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
RuntimeError: Expected floating point type for target with class probabilities, got Long
2

nn.CrossEntropyLoss expects LongTensors as the target containing class indices or (in newer PyTorch releases) FloatTensors in case you are using probabilities. Since you are one-hot encoding the tensor (I don’t know why you would need it) you should pass it as FloatTensors.

Also,

this is wrong, since nn.CrossEntropyLoss expects raw logits.

1 Like