Merge pull request #4511 from beam2d/fp16-serialize

Support deserializing into array of different dtype
chainer · Mar 26, 2018 · 9744c49 · 9744c49
2 parents 7aca074 + e79febd
commit 9744c49
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 2 deletions.
diff --git a/chainer/serializers/hdf5.py b/chainer/serializers/hdf5.py
@@ -138,7 +138,7 @@ def __call__(self, key, value):
         if isinstance(value, numpy.ndarray):
             dataset.read_direct(value)
         elif isinstance(value, cuda.ndarray):
-            value.set(numpy.asarray(dataset))
+            value.set(numpy.asarray(dataset, dtype=value.dtype))
         else:
             value = type(value)(numpy.asarray(dataset))
         return value

diff --git a/chainer/serializers/npz.py b/chainer/serializers/npz.py
@@ -148,7 +148,7 @@ def __call__(self, key, value):
         elif isinstance(value, numpy.ndarray):
             numpy.copyto(value, dataset)
         elif isinstance(value, cuda.ndarray):
-            value.set(numpy.asarray(dataset))
+            value.set(numpy.asarray(dataset, dtype=value.dtype))
         else:
             value = type(value)(numpy.asarray(dataset))
         return value

diff --git a/tests/chainer_tests/serializers_tests/test_hdf5.py b/tests/chainer_tests/serializers_tests/test_hdf5.py
@@ -153,6 +153,20 @@ def test_deserialize_none_value_gpu(self):
         y = numpy.empty((2, 3), dtype=numpy.float32)
         self.check_deserialize_none_value(cuda.to_gpu(y))
 
+    def test_deserialize_different_dtype_cpu(self):
+        y = numpy.empty((2, 3), dtype=numpy.float16)
+        ret = self.deserializer('y', y)
+        numpy.testing.assert_array_equal(y, self.data.astype(numpy.float16))
+        self.assertIs(ret, y)
+
+    @attr.gpu
+    def test_deserialize_different_dtype_gpu(self):
+        y = cuda.cupy.empty((2, 3), dtype=numpy.float16)
+        ret = self.deserializer('y', y)
+        numpy.testing.assert_array_equal(
+            y.get(), self.data.astype(numpy.float16))
+        self.assertIs(ret, y)
+
     def test_deserialize_scalar(self):
         z = 5
         ret = self.deserializer('z', z)

diff --git a/tests/chainer_tests/serializers_tests/test_npz.py b/tests/chainer_tests/serializers_tests/test_npz.py
@@ -151,6 +151,20 @@ def test_deserialize_gpu_strip_slashes(self):
         y = numpy.empty((2, 3), dtype=numpy.float32)
         self.check_deserialize(cuda.to_gpu(y), '/y')
 
+    def test_deserialize_different_dtype_cpu(self):
+        y = numpy.empty((2, 3), dtype=numpy.float16)
+        ret = self.deserializer('y', y)
+        numpy.testing.assert_array_equal(y, self.data.astype(numpy.float16))
+        self.assertIs(ret, y)
+
+    @attr.gpu
+    def test_deserialize_different_dtype_gpu(self):
+        y = cuda.cupy.empty((2, 3), dtype=numpy.float16)
+        ret = self.deserializer('y', y)
+        numpy.testing.assert_array_equal(
+            y.get(), self.data.astype(numpy.float16))
+        self.assertIs(ret, y)
+
     def test_deserialize_scalar(self):
         z = 5
         ret = self.deserializer('z', z)