From dd538adf653f7d0e6e5a771a92daf4a2898e1c01 Mon Sep 17 00:00:00 2001 From: Ningxin Date: Tue, 14 Jul 2020 09:51:17 +0000 Subject: [PATCH 1/7] Add support for relu_ Signed-off-by: Ningxin --- src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py b/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py index 47aa8087df..f55f6ec070 100644 --- a/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py +++ b/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py @@ -222,6 +222,7 @@ def __repr__(self): 'ReLU': lambda module_masks, mask: relu_inshape(module_masks, mask), 'ReLU6': lambda module_masks, mask: relu_inshape(module_masks, mask), 'aten::relu': lambda module_masks, mask: relu_inshape(module_masks, mask), + 'aten::relu_': lambda module_masks, mask: relu_inshape(module_masks, mask), 'Conv2d': lambda module_masks, mask: conv2d_inshape(module_masks, mask), 'MaxPool2d': lambda module_masks, mask: maxpool2d_inshape(module_masks, mask), 'aten::max_pool2d': lambda module_masks, mask: maxpool2d_inshape(module_masks, mask), From 2bf99c23f4043fc2c9009fe4eba4a440e56068ea Mon Sep 17 00:00:00 2001 From: Ningxin Date: Wed, 15 Jul 2020 05:23:37 +0000 Subject: [PATCH 2/7] Fix several small issues. Signed-off-by: Ningxin --- .../nni/compression/torch/speedup/infer_shape.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py b/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py index f55f6ec070..0e60d61df2 100644 --- a/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py +++ b/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py @@ -414,7 +414,8 @@ def linear_inshape(module_masks, mask): """ assert isinstance(mask, CoarseMask) assert mask.mask_index[0] is None - assert module_masks.input_mask is None + if module_masks.input_mask is not None: + assert module_masks.input_mask <= mask module_masks.set_input_mask(mask) return None @@ -452,7 +453,10 @@ def view_inshape(module_masks, mask, shape): assert mask.mask_index[0] is None assert mask.mask_index[2] is None assert mask.mask_index[3] is None - assert module_masks.input_mask is None + # due to the cat operation, the same node may be + # accessed more than once + if module_masks.input_mask is not None: + assert module_masks.input_mask <= mask module_masks.set_input_mask(mask) output_cmask = CoarseMask(num_dim=2) index = [] @@ -536,12 +540,9 @@ def relu_inshape(module_masks, mask): The mask of its output tensor """ assert isinstance(mask, CoarseMask) - # TODO: double check this assert, is it possible that a module is passed twice if module_masks.input_mask is not None: # check if has a mask conflict - assert module_masks.input_mask == mask - # No need to pass the mask again - return None + assert module_masks.input_mask <= mask # assert module_masks.input_mask is None, "A relu op can only be processed once" module_masks.set_input_mask(mask) module_masks.set_output_mask(mask) From d4e886ecf9fc682df0c5b2cdb16ec3547a6e41ff Mon Sep 17 00:00:00 2001 From: Ningxin Date: Wed, 15 Jul 2020 07:52:27 +0000 Subject: [PATCH 3/7] Fix a small issue. Signed-off-by: Ningxin --- .../nni/compression/torch/speedup/infer_shape.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py b/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py index 0e60d61df2..2635617031 100644 --- a/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py +++ b/src/sdk/pynni/nni/compression/torch/speedup/infer_shape.py @@ -242,7 +242,8 @@ def __repr__(self): 'aten::cat': lambda module_mask, mask, cat_info, last_visited: cat_inshape(module_mask, mask, cat_info, last_visited), 'aten::mean': lambda module_masks, mask, shape: mean_inshape(module_masks, mask, shape), 'Dropout': lambda module_masks, mask: dropout_inshape(module_masks, mask), - 'Dropout2d': lambda module_masks, mask: dropout_inshape(module_masks, mask) + 'Dropout2d': lambda module_masks, mask: dropout_inshape(module_masks, mask), + 'aten::dropout': lambda module_masks, mask: dropout_inshape(module_masks, mask) } """ @@ -259,8 +260,14 @@ def dropout_inshape(module_masks, mask): return module_masks.output_mask # if alreay visited assert module_masks.input_mask <= mask - if module_masks.input_mask == mask: - return None + # It should be the same, we pass the masks by the reference(not the value), + # so they acutually are two references of the same object(mask, + # module_masks.input_mask). So we should continue pass the mask + # to the following nodes even module_masks.input_mask == mask. + # if pass the mask by copy.deepcopy(), then we can stop when + # module_masks.input_mask == mask. + # if module_masks.input_mask == mask: + # return None module_masks.set_input_mask(mask) module_masks.set_output_mask(mask) return module_masks.output_mask From 532e6d3d251045c66694c5261f5512c6032aab32 Mon Sep 17 00:00:00 2001 From: Ningxin Date: Thu, 23 Jul 2020 07:27:47 +0000 Subject: [PATCH 4/7] Support the desnet and inception_v3. Signed-off-by: Ningxin --- src/sdk/pynni/tests/test_model_speedup.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/sdk/pynni/tests/test_model_speedup.py b/src/sdk/pynni/tests/test_model_speedup.py index a06f991c97..392e4fe536 100644 --- a/src/sdk/pynni/tests/test_model_speedup.py +++ b/src/sdk/pynni/tests/test_model_speedup.py @@ -145,18 +145,23 @@ def test_speedup_bigmodel(self): assert model.backbone2.fc1.in_features == int(orig_model.backbone2.fc1.in_features * SPARSITY) def test_speedup_integration(self): - for model_name in ['resnet18', 'squeezenet1_1', 'mobilenet_v2']: + for model_name in ['resnet18', 'squeezenet1_1', 'mobilenet_v2', 'densenet121', 'inception_v3']: Model = getattr(models, model_name) - net = Model(pretrained=True, progress=False).to(device) + if model_name == 'inception_v3': + # jit.trace cannot capture the aux_logits path when the net.training is False + net = Model(pretrained=True, progress=False, aux_logits=False).to(device) + speedup_model = Model(aux_logits=False).to(device) + else: + net = Model(pretrained=True, progress=False).to(device) + speedup_model = Model().to(device) net.eval() # this line is necessary + speedup_model.eval() # random generate the prune config for the pruner cfgs = generate_random_sparsity(net) pruner = L1FilterPruner(net, cfgs) pruner.compress() pruner.export_model(MODEL_FILE, MASK_FILE) pruner._unwrap_model() - speedup_model = Model().to(device) - speedup_model.eval() state_dict = torch.load(MODEL_FILE) speedup_model.load_state_dict(state_dict) zero_bn_bias(net) From 2dd818b1fe8c008e2a28b2b4681231f25d8dfcdc Mon Sep 17 00:00:00 2001 From: Ningxin Date: Fri, 31 Jul 2020 01:39:30 +0000 Subject: [PATCH 5/7] Skip the un-traced node. Signed-off-by: Ningxin --- .../pynni/nni/compression/torch/speedup/compressor.py | 7 +++++++ src/sdk/pynni/tests/test_model_speedup.py | 9 ++------- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/sdk/pynni/nni/compression/torch/speedup/compressor.py b/src/sdk/pynni/nni/compression/torch/speedup/compressor.py index b31acfe664..bf6acdd7e1 100644 --- a/src/sdk/pynni/nni/compression/torch/speedup/compressor.py +++ b/src/sdk/pynni/nni/compression/torch/speedup/compressor.py @@ -141,6 +141,13 @@ def infer_modules_masks(self): """ for module_name, mask in self.masks.items(): _logger.debug('Start mask inference from %s', module_name) + if module_name not in self.torch_graph.name_to_node: + # this module is not traced in the torch_graph, + # jit.trace only correctly records functions and + # modules which are not data dependent (e.g., do + # not have conditionals on data in tensors) + # so, if a node is not traced, we just skip it. + continue self.infer_module_mask(module_name, None, mask=mask) def replace_compressed_modules(self): diff --git a/src/sdk/pynni/tests/test_model_speedup.py b/src/sdk/pynni/tests/test_model_speedup.py index 392e4fe536..845ed793ff 100644 --- a/src/sdk/pynni/tests/test_model_speedup.py +++ b/src/sdk/pynni/tests/test_model_speedup.py @@ -147,13 +147,8 @@ def test_speedup_bigmodel(self): def test_speedup_integration(self): for model_name in ['resnet18', 'squeezenet1_1', 'mobilenet_v2', 'densenet121', 'inception_v3']: Model = getattr(models, model_name) - if model_name == 'inception_v3': - # jit.trace cannot capture the aux_logits path when the net.training is False - net = Model(pretrained=True, progress=False, aux_logits=False).to(device) - speedup_model = Model(aux_logits=False).to(device) - else: - net = Model(pretrained=True, progress=False).to(device) - speedup_model = Model().to(device) + net = Model(pretrained=True, progress=False).to(device) + speedup_model = Model().to(device) net.eval() # this line is necessary speedup_model.eval() # random generate the prune config for the pruner From 7cfc3708c1afd005f57607fba216b968b682231c Mon Sep 17 00:00:00 2001 From: Ningxin Date: Fri, 31 Jul 2020 02:26:12 +0000 Subject: [PATCH 6/7] Add warning information when skip a untraced node. Signed-off-by: Ningxin --- src/sdk/pynni/nni/compression/torch/speedup/compressor.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/sdk/pynni/nni/compression/torch/speedup/compressor.py b/src/sdk/pynni/nni/compression/torch/speedup/compressor.py index bf6acdd7e1..b999eaab78 100644 --- a/src/sdk/pynni/nni/compression/torch/speedup/compressor.py +++ b/src/sdk/pynni/nni/compression/torch/speedup/compressor.py @@ -147,6 +147,7 @@ def infer_modules_masks(self): # modules which are not data dependent (e.g., do # not have conditionals on data in tensors) # so, if a node is not traced, we just skip it. + _logger.info('Warning: %s not found in the traced graph, just skip it.', module_name) continue self.infer_module_mask(module_name, None, mask=mask) From bb39e944b7391a96fe3fdf26c40432821630144d Mon Sep 17 00:00:00 2001 From: Ningxin Date: Fri, 31 Jul 2020 02:35:02 +0000 Subject: [PATCH 7/7] update the warning information. Signed-off-by: Ningxin --- src/sdk/pynni/nni/compression/torch/speedup/compressor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sdk/pynni/nni/compression/torch/speedup/compressor.py b/src/sdk/pynni/nni/compression/torch/speedup/compressor.py index b999eaab78..41753e1c9f 100644 --- a/src/sdk/pynni/nni/compression/torch/speedup/compressor.py +++ b/src/sdk/pynni/nni/compression/torch/speedup/compressor.py @@ -147,7 +147,7 @@ def infer_modules_masks(self): # modules which are not data dependent (e.g., do # not have conditionals on data in tensors) # so, if a node is not traced, we just skip it. - _logger.info('Warning: %s not found in the traced graph, just skip it.', module_name) + _logger.warning('%s has mask, but not found in the traced graph, just skip it.', module_name) continue self.infer_module_mask(module_name, None, mask=mask)