Repository: pfnet/pytorch-pfn-extras

Issue:
Snapshot extensions can't deal with `reporting.Summary` when the tensor is on GPU
As described here, `Summay.add` can deal with a tensor on GPU.
https://github.com/pfnet/pytorch-pfn-extras/blob/5024cf746846dd41a010eedb287925ee09b17e43/pytorch_pfn_extras/reporting.py#L261-L271

When calling `add` after snapshot extensions load summary from their snapshot, it causes an error because of multiple devices.

```python
import pytorch_pfn_extras as ppe
import torch

SNAPSHOT_FILENAME = "snapshot"

tensor = torch.zeros(0, device="cuda")
summary = ppe.reporting.Summary()
summary.add(tensor)
torch.save(summary.state_dict(), SNAPSHOT_FILENAME)

manager = ppe.training.ExtensionsManager(
    {}, None, None, iters_per_epoch=1, out_dir="./"
)
snapshot = ppe.training.extensions.snapshot_object(
    summary,
    SNAPSHOT_FILENAME,
    autoload=True,
)

snapshot.initialize(manager)
summary.add(tensor)
```

Output:
```
Traceback (most recent call last):
  File "/tmp/a.py", line 21, in <module>
    summary.add(tensor)
  File "/usr/local/lib/python3.8/site-packages/pytorch_pfn_extras/reporting.py", line 272, in add
    self._x += weight * value
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/reporting.py b/pytorch_pfn_extras/reporting.py
--- a/pytorch_pfn_extras/reporting.py
+++ b/pytorch_pfn_extras/reporting.py
@@ -297,15 +297,21 @@ def make_statistics(self):
     def state_dict(self):
         state = {}
         try:
-            state = {'_x': self._x, '_x2': self._x2, '_n': self._n}
+            # Save the stats as python scalars in order to avoid
+            # different device errors when loading them back
+            state = {'_x': float(self._x),
+                     '_x2': float(self._x2),
+                     '_n': int(self._n)}
         except KeyError:
             warnings.warn('The previous statistics are not saved.')
         return state
 
     def load_state_dict(self, to_load):
-        self._x = _nograd(to_load['_x'])
-        self._x2 = _nograd(to_load['_x2'])
-        self._n = _nograd(to_load['_n'])
+        # Casting here is because of backward compatibility
+        # Restore previously taken snapshots with autoload
+        self._x = float(_nograd(to_load['_x']))
+        self._x2 = float(_nograd(to_load['_x2']))
+        self._n = int(_nograd(to_load['_n']))
 
 
 class DictSummary:


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/test_reporter.py b/tests/pytorch_pfn_extras_tests/test_reporter.py
--- a/tests/pytorch_pfn_extras_tests/test_reporter.py
+++ b/tests/pytorch_pfn_extras_tests/test_reporter.py
@@ -240,19 +240,27 @@ def _check_summary_serialize(value1, value2, value3):
     with tempfile.NamedTemporaryFile() as f:
         f.close()
         torch.save(summary.state_dict(), f.name)
-        summary2.load_state_dict(torch.load(f.name))
+        # Load tensors in CPU to simulate a snapshot restore
+        summary2.load_state_dict(
+            torch.load(f.name, map_location=torch.device('cpu')))
     summary2.add(value3)
 
-    expected_mean = (value1 + value2 + value3) / 3.
+    expected_mean = float((value1 + value2 + value3) / 3.)
     expected_std = math.sqrt(
         (value1**2 + value2**2 + value3**2) / 3. - expected_mean**2)
 
     mean = summary2.compute_mean()
-    numpy.testing.assert_allclose(mean, _nograd(expected_mean))
+    if isinstance(mean, torch.Tensor):
+        mean = mean.cpu()
+    numpy.testing.assert_allclose(mean, expected_mean)
 
     mean, std = summary2.make_statistics()
-    numpy.testing.assert_allclose(mean, _nograd(expected_mean))
-    numpy.testing.assert_allclose(std, _nograd(expected_std))
+    if isinstance(mean, torch.Tensor):
+        mean = mean.cpu()
+    if isinstance(std, torch.Tensor):
+        std = std.cpu()
+    numpy.testing.assert_allclose(mean, expected_mean)
+    numpy.testing.assert_allclose(std, expected_std)
 
 
 def test_serialize_array_float():
@@ -288,6 +296,13 @@ def test_serialize_tensor():
         torch.tensor(3.5))
 
 
+def test_serialize_tensor_cuda():
+    _check_summary_serialize(
+        torch.tensor(1.5).cuda(),
+        torch.tensor(2.0).cuda(),
+        torch.tensor(3.5).cuda())
+
+
 def test_serialize_tensor_with_grad():
     _check_summary_serialize(
         torch.tensor(1.5, requires_grad=True),