From 665204755068387b99f861a04f32ac3e684cb078 Mon Sep 17 00:00:00 2001 From: Serge Panev Date: Thu, 4 Jun 2020 09:42:31 -0700 Subject: [PATCH] Delete args from Block.params Signed-off-by: Serge Panev --- python/mxnet/gluon/block.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/python/mxnet/gluon/block.py b/python/mxnet/gluon/block.py index 8a61ba572e90..45e14977e0c4 100644 --- a/python/mxnet/gluon/block.py +++ b/python/mxnet/gluon/block.py @@ -1026,6 +1026,18 @@ def _build_cache(self, *args): for name in out.list_auxiliary_states()} # Partition the graph. out = out.optimize_for(self._backend, arg_dict, aux_dict, ctx, **self._backend_opts) + # BFS to delete the delete args/aux from the block Params and its children's Params + input_names = out.list_inputs() + queue = [self] + while len(queue) > 0: + curr_block = queue.pop(0) + curr_params_names = list(curr_block.params._params.keys()) + for k in curr_params_names: + if k not in input_names: + curr_block.params._params.pop(k) + + queue.extend(curr_block._children.values()) + #update cached graph with partitioned graph self._cached_graph = data, out