Repository: pytorch/pytorch

Issue:
Should torch.load and torch.save take a device (since PyTorch 0.4)
Reading the [PyTorch 0.4.0 migration guide](https://pytorch.org/2018/04/22/0_4_0-migration-guide.html) I came across the [`torch.device`](https://pytorch.org/docs/stable/tensor_attributes.html) abstraction.

It is great for creating a device once and then passing it around and using the `.to` functions, e.g. as in:

```python
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
inputs = data.to(device)
```

The [`torch.load`](https://github.com/pytorch/pytorch/blob/d564ecb4a515e34184c976617f57b2eb01665660/torch/serialization.py#L241) function does not take a device, though. Instead it takes a `map_location` argument which can either be a lambda, a mapping, or since https://github.com/pytorch/pytorch/pull/4203 it can be a string like `'cpu'`.

Now the question is why are there these two different concepts and can they be unified into one device abstraction? Otherwise we can pass the device around _except_ for serialization where we need to transform the device abstraction into a `map_location` parameter.

Can we unify these concepts behind an API like the following?

```python
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
restored = torch.load('model.pth', device=device)
```

Related: https://github.com/pytorch/pytorch/issues/6630 - `torch.save` should also take a device


Discussion:
Yes, this sounds like a good idea.  We also have a similar issue where torch.cuda device managers currently only take ordinals, but should be able to take devices, strings.

Resolving patch:
diff --git a/torch/serialization.py b/torch/serialization.py
--- a/torch/serialization.py
+++ b/torch/serialization.py
@@ -70,6 +70,19 @@ def _cuda_deserialize(obj, location):
             device = 0
         else:
             device = max(int(location[5:]), 0)
+
+        if not torch.cuda.is_available():
+            raise RuntimeError('Attempting to deserialize object on a CUDA '
+                               'device but torch.cuda.is_available() is False. '
+                               'If you are running on a CPU-only machine, '
+                               'please use torch.load with map_location=\'cpu\' '
+                               'to map your storages to the CPU.')
+        if device >= torch.cuda.device_count():
+            raise RuntimeError('Attempting to deserialize object on CUDA device '
+                               '{} but torch.cuda.device_count() is {}. Please use '
+                               'torch.load with map_location to map your storages '
+                               'to an existing device.'.format(
+                                   device, torch.cuda.device_count()))
         return obj.cuda(device)
 
 


Test patch:
diff --git a/test/test_cuda.py b/test/test_cuda.py
--- a/test/test_cuda.py
+++ b/test/test_cuda.py
@@ -1,3 +1,4 @@
+import io
 import math
 import tempfile
 import re
@@ -9,7 +10,7 @@
 import torch.cuda.comm as comm
 
 from test_torch import TestTorch
-from common import TestCase, get_gpu_type, to_gpu, freeze_rng_state, run_tests
+from common import TestCase, get_gpu_type, to_gpu, freeze_rng_state, run_tests, PY3
 
 HAS_CUDA = True
 if not torch.cuda.is_available():
@@ -1099,6 +1100,20 @@ def test_cat_bad_input_sizes(self):
         z = torch.randn(2, 2, 1).cuda()
         self.assertRaises(RuntimeError, lambda: torch.cat([x, y, z], dim=1))
 
+    @unittest.skipIf(torch.cuda.device_count() >= 10, "Loading a cuda:9 tensor")
+    @unittest.skipIf(not PY3, "Tensor was serialized with Python 3")
+    def test_load_nonexistent_device(self):
+        # Setup: create a serialized file object with a 'cuda:9' restore location
+        tensor = torch.randn(2, device='cuda')
+        buf = io.BytesIO()
+        torch.save(tensor, buf)
+        # NB: this might not work in the future if serialization changes
+        buf = io.BytesIO(buf.getvalue().replace(b'cuda:0', b'cuda:9'))
+
+        msg = r'Attempting to deserialize object on CUDA device 9'
+        with self.assertRaisesRegex(RuntimeError, msg):
+            _ = torch.load(buf)
+
     def test_serialization(self):
         x = torch.randn(4, 4).cuda()
         with tempfile.NamedTemporaryFile() as f:
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -6446,6 +6446,30 @@ def check_map_locations(map_locations, tensor_class, intended_device):
                 torch.device('cuda', torch.cuda.device_count() - 1)
             )
 
+    @unittest.skipIf(torch.cuda.is_available(), "Testing torch.load on CPU-only machine")
+    @unittest.skipIf(not PY3, "Test tensors were serialized using python 3")
+    def test_load_nonexistent_device(self):
+        # Setup: create a serialized file object with a 'cuda:0' restore location
+        # The following was generated by saving a torch.randn(2, device='cuda') tensor.
+        serialized = (b'\x80\x02\x8a\nl\xfc\x9cF\xf9 j\xa8P\x19.\x80\x02M\xe9'
+                      b'\x03.\x80\x02}q\x00(X\x10\x00\x00\x00protocol_versionq'
+                      b'\x01M\xe9\x03X\r\x00\x00\x00little_endianq\x02\x88X\n'
+                      b'\x00\x00\x00type_sizesq\x03}q\x04(X\x05\x00\x00\x00shortq'
+                      b'\x05K\x02X\x03\x00\x00\x00intq\x06K\x04X\x04\x00\x00\x00'
+                      b'longq\x07K\x04uu.\x80\x02ctorch._utils\n_rebuild_tensor_v2'
+                      b'\nq\x00((X\x07\x00\x00\x00storageq\x01ctorch\nFloatStorage'
+                      b'\nq\x02X\x0e\x00\x00\x0094919395964320q\x03X\x06\x00\x00'
+                      b'\x00cuda:0q\x04K\x02Ntq\x05QK\x00K\x02\x85q\x06K\x01\x85q'
+                      b'\x07\x89Ntq\x08Rq\t.\x80\x02]q\x00X\x0e\x00\x00\x00'
+                      b'94919395964320q\x01a.\x02\x00\x00\x00\x00\x00\x00\x00\xbb'
+                      b'\x1f\x82\xbe\xea\x81\xd1>')
+
+        buf = io.BytesIO(serialized)
+
+        error_msg = r'Attempting to deserialize object on a CUDA device'
+        with self.assertRaisesRegex(RuntimeError, error_msg):
+            _ = torch.load(buf)
+
     def test_serialization_filelike_api_requirements(self):
         filemock = FilelikeMock(b'', has_readinto=False)
         tensor = torch.randn(3, 5)