How maintain tensor device compatibility in custom metric

[digital-idiot]

Below is an example implementation of custom metric which inherits from torchmetrics.Metric and calculates class wise average precision:

class ClassPrecision(Metric):
    # noinspection SpellCheckingInspection
    def __init__(
            self,
            num_classes: int,
            ignore_index: int = None,
            multiclass: bool = True,
            mdmc_average: str = 'global',
            **kwargs: Any
    ):
        super(ClassPrecision, self).__init__(**kwargs)
        self.num_classes = num_classes
        self.ignore_index = ignore_index
        self.multiclass = multiclass
        self.mdmc_average = mdmc_average
        self.add_state(
            name="_count",
            default=torch.tensor(data=0.0, dtype=torch.get_default_dtype()),
            dist_reduce_fx="sum"
        )
        self.add_state(
            name='_precision',
            default=torch.zeros(
                self.num_classes,
                dtype=torch.get_default_dtype()
            ),
            dist_reduce_fx="sum"
        )

    # noinspection SpellCheckingInspection
    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
        self._count += 1.0
        self._precision += precision(
            preds=preds,
            target=target,
            average='none',
            mdmc_average=self.mdmc_average,
            num_classes=self.num_classes,
            multiclass=self.multiclass,
            ignore_index=self.ignore_index
        )

    def compute(self) -> Any:
        return self._precision / self._count

I use this in my LightningModule as:

class LightningWrapper(LightningModule):
    def __init__(self, ...)
        ...
        self.precision = ClassPrecision(...)
    
    def training_step(self, ...)
        ...
        self.precision.update(...)

However, I get:

    self._precision += precision(
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

How to make sure all the tensors in Metric are in the same device without explicitly casting the tensor to a specific device?

[SkafteNicki]

In this case I assume that the output of

precision(
            preds=preds,
            target=target,
            average='none',
            mdmc_average=self.mdmc_average,
            num_classes=self.num_classes,
            multiclass=self.multiclass,
            ignore_index=self.ignore_index
        )

is the one that is on cpu and self._precision is on the cuda device. torchmetrics.functional.precision should output an tensor on gpu if the input preds and target is also on gpu. Could you check that what you pass to

self.precision.update(...)

is on the gpu?

[digital-idiot]

@SkafteNicki I didn't explicitly cast the tensors to specific devices manually. I instead passed accelerator="gpu" when creating the Trainer instance.
From the documentation:

Modular metrics are automatically placed on the correct device when properly defined inside a LightningModule. This means that your data will always be placed on the same device as your metrics. No need to call .to(device) anymore!

So I kind of expected that API will automatically take care of the casting to appropriate device. Anyway, I later noticed that torchmetrics now have ClasswiseWrapper and MetricCollection which is pretty convenient. I will try them instead of trying to reinvent the wheel.

MetricCollection is not added to the latest documentation yet so I guess it has not been released for use yet? I tried it anyway. I implemented a metric collection using MetricCollection as given below. Notice that all the metrics have the same call signature.

class SegmentationMetrics(MetricCollection):

    def __init__(
            self,
            num_classes: int,
            multiclass: bool = True,
            ignore_index: int = None,
            prefix: str = None,
            postfix: str = None,
            normalization: str = 'true'
    ):
        metrics_dict = {
            'Precision': Precision(
                num_classes=num_classes,
                threshold=0.5,
                mdmc_average='samplewise',
                ignore_index=ignore_index,
                average='none',
                top_k=None,
                multiclass=multiclass
            ),
            'IoU': JaccardIndex(
                num_classes=num_classes,
                ignore_index=ignore_index,
                absent_score=0.0,
                threshold=0.5,
                multilabel=False,
                reduction='none'
            ),
            "Confusion_Matrix": ConfusionMatrix(
                num_classes=num_classes,
                threshold=0.5,
                multilabel=False,
                normalize=normalization
            ),
            "MCC": MatthewsCorrCoef(
                num_classes=num_classes,
                threshold=0.5
            ),
            "Kappa": CohenKappa(
                num_classes=num_classes,
                weights=None,
                threshold=0.5
            )
        }
        super(SegmentationMetrics, self).__init__(
            metrics=metrics_dict,
            prefix=prefix,
            postfix=postfix,
            compute_groups=True
        )

But when I call update on this I get RuntimeError: Long did not match Float on this line. I think this is a bug because state of a metric can be integer type (example: confmat state in JaccardIndex) which is not compatible with torch.allclose() call as it expects both operands to be of similar dtype.

Related pytorch issues: #32589, #55356

[SkafteNicki]

@digital-idiot the API should absolutely take care of anything related to this, that is at least what it does for me. That is why I wanted you to check that the input was on the correct device because it should be. I can debug this further if you can provide a reproducible example.

Regarding MetricCollection: that is not a new feature, it has been present since nearly the beginning of torchmetrics. We have dedicated space for it in the overview section in the documentation:
https://torchmetrics.readthedocs.io/en/stable/pages/overview.html#metriccollection

Regarding the last error you report: that was not something I was aware of that torch.allclose is not more robust towards difference in dtypes. I will fix that before our next release.

[digital-idiot]

Regarding MetricCollection: that is not a new feature, it has been present since nearly the beginning of torchmetrics

My bad I just checked the Wrappers and API Reference section.