Repository: Lightning-AI/pytorch-lightning
Pull request: [Metrics] [Docs] Add section about device placement
## What does this PR do?

<!--
IMPORTANT:
 We separated bug-fix PRs and feature PRs and they shall land in master and release/1.X-dev accordingly.
 By default all PR are targeted to master which is correct for bug-fixes, but need to be change for features.
 If you miss it we can still fix it for you, just ping us... :]

Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.
-->

Fixes https://github.com/PyTorchLightning/pytorch-lightning/issues/4909
Adds a section to metrics docs about device placement

## Before submitting
- [ ] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [ ] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), Pull Request section?
- [ ] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [ ] Did you make sure to update the documentation with your changes [if needed]?
- [ ] Did you write any new necessary tests [no need for typos, docs]? 
- [ ] Did you verify new and existing tests pass locally with your changes?
- [ ] If you made a notable change (that affects users), did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)?

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:  

 - [ ] Is this pull request ready for review? (if not, please submit in draft mode)
 - [ ] Check that all items from **Before submitting** are resolved
 - [ ] Make sure the title is self-explanatory and the description concisely explains the PR
 - [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified
 - [ ] **Check that target branch and milestone are aligned!**
 

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
metrics fail if inputs are on gpu
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!
```

Hints from issue discussion:
@teddykoker mind have a look?
We could move the metrics to GPUS if both provided arguments are on GPUs.
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. 
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.
> 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 
Shouldn't this be moved to gpu automatically inside a pl module though? I'll give it a test shortly...
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).
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.

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:
``` python
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]))
```
Yep I've had no issues using metrics inside pl modules. This seems like an intended behavior?

Commit message: Update docs/source/metrics.rst

Co-authored-by: Shreeyak <shreeyak.sajjan@gmail.com>

Patch to review:
diff --git a/docs/source/metrics.rst b/docs/source/metrics.rst
index 387cbc3bd7482..f99ee799e28cf 100644
--- a/docs/source/metrics.rst
+++ b/docs/source/metrics.rst
@@ -137,6 +137,56 @@ This metrics API is independent of PyTorch Lightning. Metrics can directly be us
     To change this, after initializing the metric, the method ``.persistent(mode)`` can
     be used to enable (``mode=True``) or disable (``mode=False``) this behaviour.
 
+*******************
+Metrics and devices
+*******************
+
+Metrics are simple subclasses of ``torch.nn.Module`` and their metric states behave
+similar to buffers and parameters of modules. This means that metrics states should
+be moved to the same device as the input of the metric:
+
+.. code-block:: python
+
+    import torch
+    from pytorch_lightning.metrics import Accuracy
+
+    target = torch.tensor([1, 1, 0, 0], device=torch.device("cuda", 0))
+    preds = torch.tensor([0, 1, 0, 0], device=torch.device("cuda", 0))
+
+    # Metric states are always initialized on cpu, and needs to be moved to
+    # the correct device
+    confmat = Accuracy(num_classes=2).to(torch.device("cuda", 0))
+    out = confmat(preds, target)
+    print(out.device) # cuda:0
+
+However, when **properly defined** inside an ``LightningModule``, lightning will
+automatically move the metrics to the same device as the data. By **properly defined**
+it is meant that the metric is correctly identified as a child module of the model 
+(check ``.children()`` attribute of the model). Therefore, metrics cannot be placed 
+in native python ``list`` and ``dict``, as they will not be correctly identified
+as child modules. Instead of ``list`` use ``torch.nn.ModuleList`` and instead of
+``dict`` use ``torch.nn.ModuleDict``.
+
+.. code-block::
+
+    class MyModule(pl.LightningModule):
+        def __init__(self):
+            ...
+            # valid ways metrics will be identified as child modules
+            self.metric1 = pl.metrics.Accuracy()
+            self.metric2 = torch.nn.ModuleList(pl.metrics.Accuracy())
+            self.metric3 = torch.nn.ModuleDict({'accuracy': Accuracy()})
+
+        def training_step(self, batch, batch_idx):
+            # all metrics will be on the same device as the input batch
+            data, target = batch
+            preds = self(data)
+            ...
+            val1 = self.metric1(preds, target)
+            val2 = self.metric2[0](preds, target)
+            val3 = self.metric3['accuracy'](preds, target)
+
+
 *********************
 Implementing a Metric
 *********************


Code review comments:
diff_hunk: @@ -137,6 +137,56 @@ This metrics API is independent of PyTorch Lightning. Metrics can directly be us
     To change this, after initializing the metric, the method ``.persistent(mode)`` can
     be used to enable (``mode=True``) or disable (``mode=False``) this behaviour.
 
+*******************
+Metrics and devices
+*******************
+
+Metrics are simple subclasses of ``torch.nn.Module`` and their metric states behave
+similar to buffers and parameters of modules. This means that metrics states should
+be moved to the same device as the input of the metric:
+
+.. code-block:: python
+
+    import torch
+    from pytorch_lightning.metrics import Accuracy
+
+    target = torch.tensor([1, 1, 0, 0], device=torch.device("cuda", 0))
+    preds = torch.tensor([0, 1, 0, 0], device=torch.device("cuda", 0))
+
+    # Metric states are always initialized on cpu, and needs to be moved to
+    # the correct device
+    confmat = Accuracy(num_classes=2).to(torch.device("cuda", 0))
+    out = confmat(preds, target)
+    print(out.device) # cuda:0
+
+However, when **properly defined** inside an ``LightningModule``, lightning will
+automatically move the metrics to the same device as the data. By **properly defined**
path: docs/source/metrics.rst
text: @user1:
```suggestion
automatically move the metrics to the same device as the data. Being **properly defined**
```