Repository: Lightning-AI/pytorch-lightning

Issue:
validation_epoch_end needs to return CUDA tensors
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

I'm not sure if this is expected behaviour or not, but upgrading to the latest version (from 0.8.1) caused my `validation_epoch_end` to break. It appears that a CUDA tensor is expected for the metric where before the tensor was device agnostic.

This was using sklearn's `roc_auc_score`. I haven't yet got around to testing PL's new metrics.

Feel free to close if this is expected behaviour.

### To Reproduce

This is my `validation_epoch_end`. Uncommenting `.to(avg_loss.device)` allows this to run with the dev version of PL.

This was run with `ddp`, `precision=16` using Apex AMP.

```
def validation_epoch_end(self, outputs):
        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        y_pred = torch.cat([x["y_pred"].squeeze() for x in outputs]).cpu().numpy()
        y_true = torch.cat([x["y_true"].squeeze() for x in outputs]).cpu().numpy()

        metric = torch.tensor(roc_auc_score(y_true, y_pred))  # .to(avg_loss.device)

        tensorboard_logs = {
            "loss/validation": avg_loss,
            "auc": metric,
        }
        return {"val_loss": avg_loss, "log": tensorboard_logs, "auc": metric}
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
The error message can be seen here: https://github.com/PyTorchLightning/pytorch-lightning/issues/2411#issuecomment-652171874 

#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue. 
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py
# For security purposes, please check the contents of collect_env_details.py before running it.
python collect_env_details.py
```

 - PyTorch Version (e.g., 1.0): 1.5
 - OS (e.g., Linux): Ubuntu 20.04
 - How you installed PyTorch (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version: 3.8.2
 - CUDA/cuDNN version: 10.2
 - GPU models and configuration: Dual GPU (ddp)
 - Any other relevant information: Apex AMP

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
I also have a sklearn generated metric that is only logged, with no business being on gpu. Hopefully not intended behavior?
@justusschock @SkafteNicki mind have look? :rabbit: 
@Anjum48 this is probably caused by https://github.com/PyTorchLightning/pytorch-lightning/pull/2434 in combination with your dip backend (from here: https://pytorch.org/docs/stable/distributed.html#backends) which does not support CPU tensors (are you using NCCL backend?)

@s-rog can you elaborate on your problem? Maybe in a separate issue since I'm not sure, if this is the same problem here.

cc @williamFalcon ^^
Hi @justusschock, I'm using the default PyTorch install from conda, so I believe NCCL is being used (both GPUs are in the same machine). The weird thing is that this only started happening when I upgraded from 0.8.1 to 0.8.4 so I guess something has changed in ddp?
This is not on ddp itself, but we added a often requested feature and synced the outputs across ddp nodes. And since you probably used nccl backend (which only supports gpu), this probably causes the drawback. Can you try to switch the backend of DDP? Maybe to Gloo or MPI?
I added `torch.distributed.Backend('gloo')` to the top of my training script and am seeing the same error
@justusschock I have the same issue as OP where my `validation_epoch_end` returns a cpu tensor like:
`return {'log':{'some_sklearn_metric':a_cpu_tensor}}`

Edit:
does this mean that the previous implementation was only returning/logging the rank 0 output?
> Hi @justusschock, I'm using the default PyTorch install from conda, so I believe NCCL is being used (both GPUs are in the same machine). The weird thing is that this only started happening when I upgraded from 0.8.1 to 0.8.4 so I guess something has changed in ddp?

FYI, We are testing also against Conda PyTorch distributions
> does this mean that the previous implementation was only returning/logging the rank 0 output?

I was also wondering this. I noticed my validation loss double after upgrading to 0.8.4
It is synced/reduced back to process 0 and only logged there :) DO you think it would be better to log from each process with the respective rank as a prefix?
@Anjum48 regarding your persisting error:
adding it at the top of your script doesn't help, since lightning trainer overwrites this. You have to set it in the trainer: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/trainer.py#L166
Ok I tried it in trainer, and got this error (after disabling AMP):
```
Loaded pretrained weights for efficientnet-b0
train.py:141: DtypeWarning: Columns (5) have mixed types.Specify dtype option on import or set low_memory=False.
  train_single_fold(args)
GPU available: True, used: False
TPU available: False, using: 0 TPU cores

  | Name      | Type             | Params
-----------------------------------------------
0 | critereon | CrossEntropyLoss | 0     
1 | metric    | AUROC            | 0     
2 | net       | EfficientNet     | 4 M   
Traceback (most recent call last):
  File "train.py", line 141, in <module>
    train_single_fold(args)
  File "train.py", line 65, in train_single_fold
    trainer.fit(model)
  File "/home/anjum/anaconda3/envs/kaggle/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1020, in fit
    self.run_pretrain_routine(model)
  File "/home/anjum/anaconda3/envs/kaggle/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1128, in run_pretrain_routine
    self.reset_val_dataloader(ref_model)
  File "/home/anjum/anaconda3/envs/kaggle/lib/python3.8/site-packages/pytorch_lightning/trainer/data_loading.py", line 342, in reset_val_dataloader
    self._reset_eval_dataloader(model, 'val')
  File "/home/anjum/anaconda3/envs/kaggle/lib/python3.8/site-packages/pytorch_lightning/trainer/data_loading.py", line 268, in _reset_eval_dataloader
    dataloaders = self.request_dataloader(getattr(model, f'{mode}_dataloader'))
  File "/home/anjum/anaconda3/envs/kaggle/lib/python3.8/site-packages/pytorch_lightning/trainer/data_loading.py", line 363, in request_dataloader
    dataloader = dataloader_fx()
  File "/home/anjum/PycharmProjects/kaggle/siim_isic_melanoma_classification/engine.py", line 194, in val_dataloader
    sampler = DistributedSampler(dataset)
  File "/home/anjum/anaconda3/envs/kaggle/lib/python3.8/site-packages/torch/utils/data/distributed.py", line 43, in __init__
    num_replicas = dist.get_world_size()
  File "/home/anjum/anaconda3/envs/kaggle/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 582, in get_world_size
    return _get_group_size(group)
  File "/home/anjum/anaconda3/envs/kaggle/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 196, in _get_group_size
    _check_default_pg()
  File "/home/anjum/anaconda3/envs/kaggle/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 186, in _check_default_pg
    assert _default_pg is not None, \
AssertionError: Default process group is not initialized
```

I'm not familiar with gloo so I don't know what I'm looking at here. Isn't gloo for CPU only training?
fixed on master for now. 

Yes, currently only rank 0 is used for val calculation. The problem is that reducing blows up memory and causes these kinds of issues. 

however, this will be enabled again in an upcoming PR with an object that can do this kind of syncing automatically. however we’ll need to solve this problem again, so we still need to identify what the issue is. 
> It is synced/reduced back to process 0 and only logged there :)

> Yes, currently only rank 0 is used for val calculation.

@justusschock @williamFalcon ming clarifying? these two statements seem to be contradicting, unless I'm misunderstanding something...
Originally we didn't reduce on val - it was on our todo. Then we added it, but it had weird issues such as this one. So we removed it and will test it more before adding back in.
I suppose the metric package fixes/bypasses this issue?

I just converted my vanilla sk metrics to the pl sk interface, will test soon.

Edit:
```
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/metrics/metric.py", line 145, in __call__
    return apply_to_collection(self._orig_call(*args, **kwargs), torch.Tensor,
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/metrics/converters.py", line 67, in new_func
    return func_to_apply(result, *dec_args, **dec_kwargs)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/utilities/apply_func.py", line 35, in apply_to_collection
    return function(data, *args, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/metrics/converters.py", line 252, in _sync_ddp_if_available
    async_op=False)
  File "/opt/conda/lib/python3.6/site-packages/torch/distributed/distributed_c10d.py", line 898, in all_reduce
    work = _default_pg.allreduce([tensor], opts)
RuntimeError: Tensors must be CUDA and dense
```

same issue using the built in metric interface it seems... I'm guessing the metrics need to include cuda conversion (for now) to be able to be reduced?
are you on master?

We won't include a cuda conversion here (we already have a push to self.device), since this also depends on your distributed backend (and you don't have access tp/want the computation to run on a gpu all the times.
@justusschock I just updated and tried again with the same error
do you have any CPU tensors there?
nope both inputs to the lightning sklearn metric are cuda tensors
I think I know the issue. Currently there is a problem with device propagation in `utilities.device_dtype_mixin` which causes the device property to be set only for the most outer instance. E.g. when you assign the metric to your LightningModule, the device of your metric isn't updated. So far I don't know why though.

Any Idea @Borda?
not sure at this moment, will check it...
Should this be a separate issue? as it's an issue with metrics and not validation_epoch_end specifically
Just some findings tho

Currently, classification metrics return `cpu` tensor while regression metrics return `cuda` tensor, haven't tested with sk-metrics.

[COLAB EXAMPLE](https://colab.research.google.com/drive/102qrhFRbH3Jkh28-icSkuVkgelKmxlAt?usp=sharing)
@justusschock The problem is simply that when you call .to on the nn.Module, it is not calling .to on the submodules. What the pytorch nn module does instead is only move the tensors (buffers, parameters) of these submodules, since pytorch has no such thing as a device property, the device is simply defined by the device the tensors are on. 

We can come up with a creative solution that involves going over all submodules and checking if they are instance of dtypedevicemixin and calling .to on these

But this is not a real solution because if you nest multiple levels of dtype-mixin-subclass and nn.Module in a particular way, it can still lead to the same bug

EDIT: I think I have a fix that involves using the apply function on the module, and can submit a PR. 

Resolving patch:
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -68,12 +68,6 @@ def __init__(self, *args, **kwargs):
         #: True if using amp
         self.use_amp = False
 
-        #: Current dtype
-        self._dtype = torch.float
-
-        #: device reference
-        self._device = torch.device('cpu')
-
         # optionally can be set by user
         self._example_input_array = None
 
diff --git a/pytorch_lightning/metrics/metric.py b/pytorch_lightning/metrics/metric.py
--- a/pytorch_lightning/metrics/metric.py
+++ b/pytorch_lightning/metrics/metric.py
@@ -27,8 +27,6 @@ def __init__(self, name: str):
         """
         super().__init__()
         self.name = name
-        self._dtype = torch.get_default_dtype()
-        self._device = torch.device('cpu')
 
     @abstractmethod
     def forward(self, *args, **kwargs) -> torch.Tensor:
diff --git a/pytorch_lightning/utilities/device_dtype_mixin.py b/pytorch_lightning/utilities/device_dtype_mixin.py
--- a/pytorch_lightning/utilities/device_dtype_mixin.py
+++ b/pytorch_lightning/utilities/device_dtype_mixin.py
@@ -5,8 +5,11 @@
 
 
 class DeviceDtypeModuleMixin(Module):
-    _device: ...
-    _dtype: Union[str, torch.dtype]
+
+    def __init__(self):
+        super().__init__()
+        self._dtype = torch.get_default_dtype()
+        self._device = torch.device('cpu')
 
     @property
     def dtype(self) -> Union[str, torch.dtype]:
@@ -79,17 +82,14 @@ def to(self, *args, **kwargs) -> Module:
             ExampleModule()
             >>> module.weight #doctest: +ELLIPSIS
             tensor([[...]], dtype=torch.float16)
+            >>> module.device
+            device(type='cpu')
+            >>> module.dtype
+            torch.float16
         """
         # there is diff nb vars in PT 1.5
         out = torch._C._nn._parse_to(*args, **kwargs)
-        device = out[0]
-        dtype = out[1]
-        if device is not None:
-            self._device = device
-
-        if dtype is not None:
-            self._dtype = dtype
-
+        self.__update_properties(device=out[0], dtype=out[1])
         return super().to(*args, **kwargs)
 
     def cuda(self, device: Optional[int] = None) -> Module:
@@ -105,8 +105,7 @@ def cuda(self, device: Optional[int] = None) -> Module:
         Returns:
             Module: self
         """
-
-        self._device = torch.device('cuda', index=device)
+        self.__update_properties(device=torch.device('cuda', index=device))
         return super().cuda(device=device)
 
     def cpu(self) -> Module:
@@ -114,7 +113,7 @@ def cpu(self) -> Module:
         Returns:
             Module: self
         """
-        self._device = torch.device('cpu')
+        self.__update_properties(device=torch.device('cpu'))
         return super().cpu()
 
     def type(self, dst_type: Union[str, torch.dtype]) -> Module:
@@ -126,7 +125,7 @@ def type(self, dst_type: Union[str, torch.dtype]) -> Module:
         Returns:
             Module: self
         """
-        self._dtype = dst_type
+        self.__update_properties(dtype=dst_type)
         return super().type(dst_type=dst_type)
 
     def float(self) -> Module:
@@ -135,7 +134,7 @@ def float(self) -> Module:
         Returns:
             Module: self
         """
-        self._dtype = torch.float
+        self.__update_properties(dtype=torch.float)
         return super().float()
 
     def double(self) -> Module:
@@ -144,7 +143,7 @@ def double(self) -> Module:
         Returns:
             Module: self
         """
-        self._dtype = torch.double
+        self.__update_properties(dtype=torch.double)
         return super().double()
 
     def half(self) -> Module:
@@ -153,5 +152,17 @@ def half(self) -> Module:
         Returns:
             Module: self
         """
-        self._dtype = torch.half
+        self.__update_properties(dtype=torch.half)
         return super().half()
+
+    def __update_properties(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+
+        def apply_fn(module):
+            if not isinstance(module, DeviceDtypeModuleMixin):
+                return
+            if device is not None:
+                module._device = device
+            if dtype is not None:
+                module._dtype = dtype
+
+        self.apply(apply_fn)


Test patch:
diff --git a/tests/utilities/test_dtype_device_mixin.py b/tests/utilities/test_dtype_device_mixin.py
new file mode 100644
--- /dev/null
+++ b/tests/utilities/test_dtype_device_mixin.py
@@ -0,0 +1,89 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from pytorch_lightning import Trainer, Callback
+from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
+from tests.base import EvalModelTemplate
+
+
+class SubSubModule(DeviceDtypeModuleMixin):
+    pass
+
+
+class SubModule(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.module = SubSubModule()
+
+
+class TopModule(EvalModelTemplate):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.module = SubModule()
+
+
+class DeviceAssertCallback(Callback):
+
+    def on_batch_start(self, trainer, model):
+        rank = trainer.local_rank
+        assert isinstance(model, TopModule)
+        # index = None also means first device
+        assert (model.device.index is None and rank == 0) or model.device.index == rank
+        assert model.device == model.module.module.device
+
+
+@pytest.mark.parametrize(['dst_dtype'], [
+    pytest.param(torch.float),
+    pytest.param(torch.double),
+    pytest.param(torch.half),
+])
+@pytest.mark.parametrize(['dst_device'], [
+    pytest.param(torch.device('cpu')),
+    pytest.param(torch.device('cuda')),
+    pytest.param(torch.device('cuda', 0)),
+])
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+def test_submodules_device_and_dtype(dst_device, dst_dtype):
+    """
+    Test that the device and dtype property updates propagate through mixed nesting of regular
+    nn.Modules and the special modules of type DeviceDtypeModuleMixin (e.g. Metric or LightningModule).
+    """
+
+    model = TopModule()
+    assert model.device == torch.device('cpu')
+    model = model.to(device=dst_device, dtype=dst_dtype)
+    # nn.Module does not have these attributes
+    assert not hasattr(model.module, '_device')
+    assert not hasattr(model.module, '_dtype')
+    # device and dtype change should propagate down into all children
+    assert model.device == model.module.module.device == dst_device
+    assert model.dtype == model.module.module.dtype == dst_dtype
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_submodules_multi_gpu_dp(tmpdir):
+    model = TopModule()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        distributed_backend='dp',
+        gpus=2,
+        callbacks=[DeviceAssertCallback()],
+        max_steps=1,
+    )
+    trainer.fit(model)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_submodules_multi_gpu_ddp_spawn(tmpdir):
+    model = TopModule()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        distributed_backend='dpp_spawn',
+        gpus=2,
+        callbacks=[DeviceAssertCallback()],
+        max_steps=1,
+    )
+    trainer.fit(model)