metrics fail if inputs are on gpu

[simonm3]

The metrics precision and recall fail if inputs are on gpu.

>>> from pytorch_lightning.metrics import Precision
>>> target = torch.tensor([0, 1, 2, 0, 1, 2]).cuda()
>>> preds = torch.tensor([0, 2, 1, 0, 0, 1]).cuda()
>>> precision = Precision(num_classes=3)
>>> precision(preds, target)
tensor(0.3333)

RuntimeError                              Traceback (most recent call last)
<ipython-input-39-834a61ec6e2d> in <module>
      3 preds = torch.tensor([0, 2, 1, 0, 0, 1]).cuda()
      4 precision = Precision(num_classes=3)
----> 5 precision(preds, target)
      6 tensor(0.3333)

/anaconda/envs/azureml_py36/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/anaconda/envs/azureml_py36/lib/python3.6/site-packages/pytorch_lightning/metrics/metric.py in forward(self, *args, **kwargs)
    154         # add current step
    155         with torch.no_grad():
--> 156             self.update(*args, **kwargs)
    157         self._forward_cache = None
    158 

/anaconda/envs/azureml_py36/lib/python3.6/site-packages/pytorch_lightning/metrics/metric.py in wrapped_func(*args, **kwargs)
    200         def wrapped_func(*args, **kwargs):
    201             self._computed = None
--> 202             return update(*args, **kwargs)
    203         return wrapped_func
    204 

/anaconda/envs/azureml_py36/lib/python3.6/site-packages/pytorch_lightning/metrics/classification/precision_recall.py in update(self, preds, target)
    130 
    131         # multiply because we are counting (1, 1) pair for true positives
--> 132         self.true_positives += torch.sum(preds * target, dim=1)
    133         self.predicted_positives += torch.sum(preds, dim=1)
    134 

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

[Borda]

@teddykoker mind have a look?

[edenlightning]

We could move the metrics to GPUS if both provided arguments are on GPUs.

[SkafteNicki]

Metrics are subclass of nn.Module and their states behave similar to the parameters and buffers of other modules. Therefore, calling precision.cuda() before evaluating the metric should solve your problem.

[simonm3]

It is cleaner to be able to use the same code on cpu or gpu without having multiple places where you have to test and convert.

[Borda]

It is cleaner to be able to use the same code on cpu or gpu without having multiple places where you have to test and convert.

I agree here, this would be very useful, basically follow the agnostic philosophy as Trainer does...
cc: @PyTorchLightning/core-contributors

[s-rog]

Shouldn't this be moved to gpu automatically inside a pl module though? I'll give it a test shortly...

[SkafteNicki]

If the metric is initialized inside a LightningModule or a nn.Module it will automatically be moved to the correct device, whenever the .to(...) is called. However, if used outside a model the user would need to transfer the metric to the correct device. This is completely inline with other pytorch modules, that also needs to be moved to correct device by the user (take torch.nn.Conv2d for example).

[simonm3]

Not sure what you mean by inside a model. I have:

class Module(pl.LightningModule):
    def __init__(self):
        super().__init__()
        model = resnet18(pretrained=True)
        self.metrics = dict(precision=Precision())
        self.loss = partial(F.cross_entropy, weight=torch.tensor([0.1, 0.9]))

In training_step and validation_step I call self.loss and each of the self.metrics. precision(y,ypred) fails if ypred and y are both gpu. so I have to force them to cpu.

The loss also fails unless I explicitly change the above to weight=torch.tensor([0.1, 0.9]).cuda(). Everything else seems to detect whether it is running on cpu or gpu.

[SkafteNicki]

Then the problem has to with how pytorch register child modules. Essentially, anything put inside a list or dict will not be registered correctly. Instead you should be using nn.ModuleList and nn.ModuleDict. Addtionally, you can use the module version of the loss, and the weight should also automatically transferred to the correct device. This version of your code should work:

class Module(pl.LightningModule):
    def __init__(self):
        super().__init__()
        model = resnet18(pretrained=True)
        self.metrics = torch.nn.ModuleDict({'precision' : Precision()})
        self.loss = torch.nn.CrossEntropyLoss(weight=torch.tensor([0.1, 0.9]))

[s-rog]

Yep I've had no issues using metrics inside pl modules. This seems like an intended behavior?

[awaelchli]

agree with @SkafteNicki this is most likely the reason. that's a silly mistake every pytorch user encounters at least once in their life xD
one just has to know this :)