Best place to .to(device) a torchmetrics.Metric

[ZhaofengWu]

In the init of my LightningModule, I have something like

def __init__(self):
  super().init()
  ...
  self.metrics = {"dev": {"accuracy": Accuracy(...), ...some other metrics}, ...some other splits}

I'm wondering where's the best place to move these torchmetrisc.Metrics to GPU? Should I override def to?

[ZhaofengWu]

Bump

[rohitgr7]

can you explain more on what do you mean by moving the metrics to the device. What do you want to do exactly?

[ZhaofengWu]

So I think by default, if I do something like metric(predictions, labels) where predictions and labels are GPU tensors, torchmetrics throws a device mismatch error, I think because the internal state tracking tensors are on CPU. torchmetrics provides a .to(device) method that moves the metric to GPU, getting rid of this error, but I'm unsure what's the best way to do that is in the lightning framework. Alternatively, I could move predictions and labels to CPU, but I think some metric computations could benefit from GPU speedup.

[rohitgr7]

I don't think it moves the tensors to CPU. But if you facing this error can you tell which metric you are using here? or maybe share a small code-snippet to check this?

[ZhaofengWu]

With torchmetrics==0.10.0:

Python 3.10.6 (main, Oct  7 2022, 20:19:58) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> from torchmetrics import Accuracy
>>> acc = Accuracy(num_classes=2)
>>> acc(torch.tensor([[1., 2.], [3., 4.]]).cuda(), torch.tensor([1, 1]).cuda())
Traceback (most recent call last):
  File ".../torchmetrics/metric.py", line 395, in wrapped_func
    update(*args, **kwargs)
  File ".../torchmetrics/classification/accuracy.py", line 604, in update
    self.tp += tp
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File ".../torchmetrics/metric.py", line 245, in forward
    self._forward_cache = self._forward_reduce_state_update(*args, **kwargs)
  File ".../torchmetrics/metric.py", line 309, in _forward_reduce_state_update
    self.update(*args, **kwargs)
  File ".../torchmetrics/metric.py", line 398, in wrapped_func
    raise RuntimeError(
RuntimeError: Encountered different devices in metric calculation (see stacktrace for details). This could be due to the metric class not being on the same device as input. Instead of `metric=Accuracy(...)` try to do `metric=Accuracy(...).to(device)` where device corresponds to the device of the input.

[rohitgr7]

oh I see...
didn't know there's such a thing. For this I'd suggest maybe inside:

def on_fit_start(self):
    self.acc = self.acc.to(self.device)

def on_validation_start(self)
    self.val_acc = self.val_acc.to(self.device)

def on_test_start(self):  # if you are using `trainer.test` and doing test logging
    self.test_acc = self.test_acc.to(self.device)

id you found this on docs somewhere, can you share that? we should update that with this example :)

[ZhaofengWu]

Thanks! This makes sense. Alternatively, would it make sense to override def to instead? Semantically that sounds like the right thing to do (and is shorter), but it seems like lightning doesn't even override it, so I'd be overriding pytorch's method directly, which sounds might be a bit dangerous. Regardless, given the close affinity between lightning and torchmetrics, ideally there could be some official doc on this. Is this also the doc you're asking about, or something else?