Repository: pytorch/pytorch

Issue:
expose CPU HalfTensor, fix GPU HalfTensor serialization
An issue to track exposing CPU half tensors.
Note: CPU half tensors don't have math

cc: @ajbrock

Right now, calling `.cpu()` on a `torch.cuda.HalfTensor` gives an error.
```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-8-29bf7ba9ae5b> in <module>()
----> 1 c.cpu()

/home/fatbox/miniconda2/lib/python2.7/site-packages/torch/tensor.pyc in cpu(self)
     48     def cpu(self):
     49         """Returns a CPU copy of this tensor if it's not already on the CPU"""
---> 50         return self.type(getattr(torch, self.__class__.__name__))
     51
     52     def double(self):

AttributeError: 'module' object has no attribute 'HalfTensor'
```

Also, saving `torch.cuda.HalfTensor` with `torch.save` gives the error:

```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-9-ed138460cf2d> in <module>()
----> 1 torch.save(c, 'c.pth')

/home/fatbox/miniconda2/lib/python2.7/site-packages/torch/serialization.pyc in save(obj, f, pickle_module, pickle_protocol)
    121         f = open(f, "wb")
    122     try:
--> 123         return _save(obj, f, pickle_module, pickle_protocol)
    124     finally:
    125         if new_fd:

/home/fatbox/miniconda2/lib/python2.7/site-packages/torch/serialization.pyc in _save(obj, f, pickle_module, pickle_protocol)
    216         _add_to_tar(pickle_objects, tar, 'pickle')
    217         _add_to_tar(save_tensors, tar, 'tensors')
--> 218         _add_to_tar(save_storages, tar, 'storages')
    219
    220

/home/fatbox/miniconda2/lib/python2.7/site-packages/torch/serialization.pyc in _add_to_tar(fn, tar_file, name)
     29 def _add_to_tar(fn, tar_file, name):
     30     tmp_file = tempfile.NamedTemporaryFile(delete=False)
---> 31     fn(tmp_file)
     32     tmp_file.close()
     33

/home/fatbox/miniconda2/lib/python2.7/site-packages/torch/serialization.pyc in save_storages(f)
    187         for key, storage in serialized_storages.items():
    188             location = location_tag(storage)
--> 189             storage_type = normalize_storage_type(type(storage))
    190             pickle_module.dump((key, location, storage_type), f,
    191                     protocol=pickle_protocol)

/home/fatbox/miniconda2/lib/python2.7/site-packages/torch/serialization.pyc in normalize_storage_type(storage_type)
     97
     98 def normalize_storage_type(storage_type):
---> 99     return getattr(torch, storage_type.__name__)
    100
    101

AttributeError: 'module' object has no attribute 'HalfStorage'
```






Discussion:

Resolving patch:
diff --git a/tools/cwrap/plugins/THPPlugin.py b/tools/cwrap/plugins/THPPlugin.py
--- a/tools/cwrap/plugins/THPPlugin.py
+++ b/tools/cwrap/plugins/THPPlugin.py
@@ -303,8 +303,6 @@ def generate_out_options(self, declaration):
 
     def process_declarations(self, declarations):
         new_declarations = []
-        register_only = [d for d in declarations if d.get('only_register', False)]
-        declarations = [d for d in declarations if not d.get('only_register', False)]
 
         def has_arg_type(declaration, type_name):
             return any(arg['type'] == type_name
@@ -322,8 +320,16 @@ def has_output_args(declaration):
                        for arg in option['arguments'])
 
         for declaration in declarations:
+            # Disable all methods for THHalfTensor, unless cpu_half is True
+            if not declaration.get('cpu_half', False):
+                defined_if = '!defined(TH_REAL_IS_HALF)'
+                if 'defined_if' in declaration:
+                    defined_if += ' && (' + declaration['defined_if'] + ')'
+                declaration['defined_if'] = defined_if
+
             if declaration.get('only_register', False):
                 continue
+
             declaration.setdefault('python_name', declaration['name'])
             declaration.setdefault('variables', [])
             if has_arg_type(declaration, 'THSize*'):
@@ -353,7 +359,9 @@ def has_output_args(declaration):
                     if arg['name'] == 'self':
                         arg['ignore_check'] = True
 
-        declarations = [d for d in declarations if not d.get('only_stateless', False)]
+        register_only = [d for d in declarations if d.get('only_register', False)]
+        declarations = [d for d in declarations
+                        if (not d.get('only_stateless', False)) and (not d.get('only_register', False))]
         self.declarations.extend(filter(lambda x: not x.get('only_stateless', False), register_only))
         self.stateless_declarations.extend(filter(lambda x: x.get('only_stateless', False), register_only))
 
@@ -390,11 +398,14 @@ def declare_methods(self, stateless, sparse):
             if 'defined_if' in declaration:
                 entry = self.preprocessor_guard(entry, declaration['defined_if'])
             tensor_methods += entry
-        return self.TENSOR_METHODS_DECLARATION.substitute(
+        generated = self.TENSOR_METHODS_DECLARATION.substitute(
             methods=tensor_methods,
             stateless=('' if not stateless else 'stateless_'),
             sparse=('' if not sparse else 'S'),
         )
+        if sparse:
+            generated = '#ifndef TH_REAL_IS_HALF\n' + generated + '\n#endif\n\n'
+        return generated
 
     def process_full_file(self, code):
         # We have to find a place before all undefs
diff --git a/torch/__init__.py b/torch/__init__.py
--- a/torch/__init__.py
+++ b/torch/__init__.py
@@ -151,6 +151,10 @@ class FloatStorage(_C.FloatStorageBase, _StorageBase):
     pass
 
 
+class HalfStorage(_C.HalfStorageBase, _StorageBase):
+    pass
+
+
 class LongStorage(_C.LongStorageBase, _StorageBase):
     pass
 
@@ -191,6 +195,16 @@ def storage_type(cls):
         return FloatStorage
 
 
+class HalfTensor(_C.HalfTensorBase, _TensorBase):
+
+    def is_signed(self):
+        return True
+
+    @classmethod
+    def storage_type(cls):
+        return HalfStorage
+
+
 class LongTensor(_C.LongTensorBase, _TensorBase):
 
     def is_signed(self):


Test patch:
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -2577,6 +2577,35 @@ def test_serialization(self):
             rootview = c[8]
             self.assertEqual(rootview.data_ptr(), c[0].data_ptr())
 
+    def test_half_tensor(self):
+        x = torch.randn(5, 5).float()
+        y = torch.randn(5, 5).float()
+        xh, yh = x.half(), y.half()
+
+        self.assertEqual(x.half().float(), x, 1e-3)
+
+        z = torch.Tensor(5, 5)
+        self.assertEqual(z.copy_(xh), x, 1e-3)
+
+        with tempfile.NamedTemporaryFile() as f:
+            torch.save(xh, f)
+            f.seek(0)
+            xh2 = torch.load(f)
+            self.assertEqual(xh, xh2)
+
+    @unittest.skipIf(not torch.cuda.is_available(), 'no CUDA')
+    def test_half_tensor_cuda(self):
+        x = torch.randn(5, 5).half()
+        self.assertEqual(x.cuda().cpu(), x)
+
+        xc = x.cuda()
+        with tempfile.NamedTemporaryFile() as f:
+            torch.save(xc, f)
+            f.seek(0)
+            xc2 = torch.load(f)
+            self.assertIsInstance(xc2, type(xc))
+            self.assertEqual(xc, xc2)
+
     @unittest.skipIf(not torch.cuda.is_available(), 'no CUDA')
     def test_serialization_cuda(self):
         device_count = torch.cuda.device_count()