Merge pull request #54 from RedisAI/tensorget_as_mutable_numpy

TESNORGET as mutable numpy ndarray

Author: alonre24

redisai/client.py

@@ -335,7 +335,7 @@ def tensorset(self,
         return res if not self.enable_postprocess else processor.tensorset(res)
 
     def tensorget(self,
-                  key: AnyStr, as_numpy: bool = True,
+                  key: AnyStr, as_numpy: bool = True, as_numpy_mutable: bool = False,
                   meta_only: bool = False) -> Union[dict, np.ndarray]:
         """
         Retrieve the value of a tensor from the server. By default it returns the numpy
@@ -349,6 +349,9 @@ def tensorget(self,
             If True, returns a numpy.ndarray. Returns the value as a list and the
             metadata in a dictionary if False. This flag also decides how to fetch
             the value from the RedisAI server, which also has performance implications
+        as_numpy_mutable : bool
+            If True, returns a a mutable numpy.ndarray object by copy the tensor data. Otherwise (as long as_numpy=True)
+            the returned numpy.ndarray will use the original tensor buffer and will be for read-only
         meta_only : bool
             If True, the value is not retrieved, only the shape and the type
 
@@ -368,8 +371,7 @@ def tensorget(self,
         """
         args = builder.tensorget(key, as_numpy, meta_only)
         res = self.execute_command(*args)
-        return res if not self.enable_postprocess else processor.tensorget(res,
-                                                                    as_numpy, meta_only)
+        return res if not self.enable_postprocess else processor.tensorget(res, as_numpy, as_numpy_mutable, meta_only)
 
     def scriptset(self, key: AnyStr, device: str, script: str, tag: AnyStr = None) -> str:
         """
@@ -587,11 +589,12 @@ def __init__(self, enable_postprocess, *args, **kwargs):
     def dag(self, *args, **kwargs):
         raise RuntimeError("Pipeline object doesn't allow DAG creation currently")
 
-    def tensorget(self, key, as_numpy=True, meta_only=False):
+    def tensorget(self, key, as_numpy=True, as_numpy_mutable=False, meta_only=False):
         self.tensorget_processors.append(partial(processor.tensorget,
                                                  as_numpy=as_numpy,
+                                                 as_numpy_mutable=as_numpy_mutable,
                                                  meta_only=meta_only))
-        return super().tensorget(key, as_numpy, meta_only)
+        return super().tensorget(key, as_numpy, as_numpy_mutable, meta_only)
 
     def _execute_transaction(self, *args, **kwargs):
         # TODO: Blocking commands like MODELRUN, SCRIPTRUN and DAGRUN won't work
@@ -648,13 +651,14 @@ def tensorset(self,
         return self
 
     def tensorget(self,
-                  key: AnyStr, as_numpy: bool = True,
+                  key: AnyStr, as_numpy: bool = True, as_numpy_mutable: bool = False,
                   meta_only: bool = False) -> Any:
-        args = builder.tensorget(key, as_numpy, meta_only)
+        args = builder.tensorget(key, as_numpy, as_numpy_mutable)
         self.commands.extend(args)
         self.commands.append("|>")
         self.result_processors.append(partial(processor.tensorget,
                                               as_numpy=as_numpy,
+                                              as_numpy_mutable=as_numpy_mutable,
                                               meta_only=meta_only))
         return self
 

redisai/postprocessor.py

@@ -19,17 +19,19 @@ def modelscan(res):
         return utils.recursive_bytetransform(res, lambda x: x.decode())
 
     @staticmethod
-    def tensorget(res, as_numpy, meta_only):
+    def tensorget(res, as_numpy, as_numpy_mutable, meta_only):
         """Process the tensorget output.
 
         If ``as_numpy`` is True, it'll be converted to a numpy array. The required
         information such as datatype and shape must be in ``rai_result`` itself.
         """
         rai_result = utils.list2dict(res)
-        if meta_only:
+        if meta_only is True:
             return rai_result
+        elif as_numpy_mutable is True:
+            return utils.blob2numpy(rai_result['blob'], rai_result['shape'], rai_result['dtype'], mutable=True)
         elif as_numpy is True:
-            return utils.blob2numpy(rai_result['blob'], rai_result['shape'], rai_result['dtype'])
+            return utils.blob2numpy(rai_result['blob'], rai_result['shape'], rai_result['dtype'], mutable=False)
         else:
             target = float if rai_result['dtype'] in ('FLOAT', 'DOUBLE') else int
             utils.recursive_bytetransform(rai_result['values'], target)

redisai/utils.py

@@ -31,14 +31,17 @@ def numpy2blob(tensor: np.ndarray) -> tuple:
     return dtype, shape, blob
 
 
-def blob2numpy(value: ByteString, shape: Union[list, tuple], dtype: str) -> np.ndarray:
+def blob2numpy(value: ByteString, shape: Union[list, tuple], dtype: str, mutable: bool) -> np.ndarray:
     """Convert `BLOB` result from RedisAI to `np.ndarray`."""
     mm = {
         'FLOAT': 'float32',
         'DOUBLE': 'float64'
     }
     dtype = mm.get(dtype, dtype.lower())
-    a = np.frombuffer(value, dtype=dtype)
+    if mutable:
+        a = np.fromstring(value, dtype=dtype)
+    else:
+        a = np.frombuffer(value, dtype=dtype)
     return a.reshape(shape)
 
 

test/test.py

@@ -102,6 +102,13 @@ def test_numpy_tensor(self):
         ret = con.tensorset('x', values)
         self.assertEqual(ret, 'OK')
 
+        # By default tensorget returns immutable, unless as_numpy_mutable is set as True
+        ret = con.tensorget('x')
+        self.assertRaises(ValueError, np.put, ret, 0, 1)
+        ret = con.tensorget('x', as_numpy_mutable=True)
+        np.put(ret, 0, 1)
+        self.assertEqual(ret[0], 1)
+
         stringarr = np.array('dummy')
         with self.assertRaises(TypeError):
             con.tensorset('trying', stringarr)