BigEarthNet - RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

[robmarkcole]

Description

Relates to Lightning-AI/pytorch-lightning#20456 I believe. Not sure if there is a workaround?

Steps to reproduce

from torchgeo.datasets import BigEarthNet
from torchgeo.datamodules import BigEarthNetDataModule
from torchgeo.trainers import ClassificationTask
from lightning.pytorch import Trainer

# Load the dataset
train_dataset = BigEarthNet(root="data", download=False) # already downloaded

# Setup datamodule
datamodule = BigEarthNetDataModule(bands='all', batch_size=16)
datamodule.setup('fit')
datamodule.setup('test')

# Define the model
num_bands = len(datamodule.maxs)
num_classes = 19  # Default number of classes

task = ClassificationTask(
    weights=True,
    num_classes=num_classes,
    in_channels=num_bands,
    lr=1e-3
)

# Train the model
trainer = Trainer(max_epochs=10)
trainer.fit(task, datamodule)

Returns

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/functional.py:3059, in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3057 if size_average is not None or reduce is not None:
   3058     reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3059 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

Version

0.6.2

[robmarkcole]

Tried setting target = target.float() and get

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

[isaaccorley]

This error is because BigEarthNet is a multilabel classification dataset but you're using ClassificationTask which is for multiclass problems. Changing ClassificationTask -> MultiLabelClassificationTask should do the trick.

[robmarkcole]

Switching to MultiLabelClassificationTask and I am getting RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

Resolved by setting:

        x = batch['image'].to(self.device)

[adamjstewart]

Yep, that's definitely Lightning-AI/pytorch-lightning#20456.