Partitioning Gluon HybridBlocks

python/mxnet/gluon/block.py

@@ -890,6 +890,8 @@ def __init__(self, prefix=None, params=None):
         self._flags = []
         self._callback = None
         self._monitor_all = False
+        self._backend = None
+        self._backend_args = {}
 
     def __setattr__(self, name, value):
         """Registers parameters."""
@@ -935,7 +937,6 @@ def _build_cache(self, *args):
         data_names = {data.name: i for i, data in enumerate(data)}
         params = self.collect_params()
         input_names = out.list_inputs()
-
         param_names = set(params.keys())
         expected_names = set(input_names)
         for name in expected_names:
@@ -967,6 +968,26 @@ def _build_cache(self, *args):
                 self._cached_op_args.append((False, params[name]))
         flags = [('data_indices', data_indices), ('param_indices', param_indices)] + \
                 self._flags
+
+        args, _ = _flatten(args, "input")
+        try:
+            for is_arg, i in self._cached_op_args:
+                if not is_arg:
+                    i.data()
+        except DeferredInitializationError:
+            self._deferred_infer_shape(*args)
+            for is_arg, i in self._cached_op_args:
+                if not is_arg:
+                    i._finish_deferred_init()
+
+        if self._backend:
+            ctx = args[0].context
+            # get list of params in the order of out.list_arguments
+            arg_array = [args[data_names[name]] if name in data_names.keys() else params[name].data()
+                         for name in out.list_arguments()]
+            # Partition the graph.
+            out = out.optimize_for(self._backend, arg_array, ctx, **self._backend_args)
+
         self._cached_op = ndarray.CachedOp(out, flags)
 
     def _deferred_infer_shape(self, *args):
@@ -1008,19 +1029,10 @@ def _call_cached_op(self, *args):
                 raise ValueError("The argument structure of HybridBlock does not match"
                                  " the cached version. Stored format = {}, input format = {}"
                                  .format(fmt, self._in_format))
+
         args_without_none = [ele for ele in args if ele is not None]
-        try:
-            cargs = [args_without_none[i] if is_arg else i.data()
-                     for is_arg, i in self._cached_op_args]
-        except DeferredInitializationError:
-            self._deferred_infer_shape(*args)
-            cargs = []
-            for is_arg, i in self._cached_op_args:
-                if is_arg:
-                    cargs.append(args_without_none[i])
-                else:
-                    i._finish_deferred_init()
-                    cargs.append(i.data())
+        cargs = [args_without_none[i] if is_arg else i.data()
+                 for is_arg, i in self._cached_op_args]
         out = self._cached_op(*cargs)
         if isinstance(out, NDArray):
             out = [out]
@@ -1040,7 +1052,12 @@ def register_child(self, block, name=None):
         super(HybridBlock, self).register_child(block, name)
         self._clear_cached_op()
 
-    def hybridize(self, active=True, **kwargs):
+    def hybridize(self, active=True, backend=None, backend_args=None, **kwargs):
+        self._backend = backend
+        if backend_args is None:
+            self._backend_args = {}
+        else:
+            self._backend_args = backend_args
         self._active = active
         self._flags = list(kwargs.items())
         self._clear_cached_op()
@@ -1160,7 +1177,6 @@ def forward(self, x, *args):
                     params = {k: v.data(ctx) for k, v in self._reg_params.items()}
 
                 return self.hybrid_forward(ndarray, x, *args, **params)
-
         params = {i: j.var() for i, j in self._reg_params.items()}
         with self.name_scope():
             return self.hybrid_forward(symbol, x, *args, **params)