fix mse strategy for mxnet (#1110)

Author: RafLit

neural_compressor/adaptor/mxnet.py

@@ -320,7 +320,8 @@ def b_filter():
         return collector.tensors_dicts
 
     def inspect_tensor(self, nc_model, data_x, op_list=[], iteration_list=[],
-                       inspect_type='activation', save_to_disk=False):
+                       inspect_type='activation', save_to_disk=False,
+                       save_path = None, quantization_cfg = None):
         """The function is used by tune strategy class for dumping tensor info.
 
         Args:
@@ -352,7 +353,7 @@ def inspect_tensor(self, nc_model, data_x, op_list=[], iteration_list=[],
                     tensor_dict[key][tensor_name] = tensor_dict[key][tensor_name].asnumpy()
 
                 # transform to format expected by neural_compressor (assume only 1 tensor for now)
-                node, op = key
+                node = key
                 assert len(tensors) == 1, 'Multiple tensors from a single node are not supported'
                 tensor = list(tensor_dict[key].values())[0]
                 tensor_dict[key] = {node: tensor}

neural_compressor/adaptor/mxnet_utils/util.py

@@ -758,11 +758,11 @@ class TensorCollector(CollectorBase):
 
     def __init__(self, include_nodes, qtensor_to_tensor, tensor_to_node):
         self.tensors_dicts = []
-        self.include_nodes = dict(include_nodes)
+        self.include_nodes = include_nodes
         self.qtensor_to_tensor = qtensor_to_tensor
         self.tensor_to_node = tensor_to_node
 
-        rest = set(self.include_nodes.keys()) - set(self.tensor_to_node.values())
+        rest = set(self.include_nodes) - set(self.tensor_to_node.values())
         assert len(rest) == 0, 'Unexpected tensors set to be collected: {}'.format(rest)
 
     def collect_gluon(self, name, _, arr):
@@ -779,9 +779,7 @@ def collect_gluon(self, name, _, arr):
 
         node = self.tensor_to_node[name]
         if node in self.include_nodes:
-            op = self.include_nodes[node]
-            key = (node, op)
-            self.tensors_dicts[-1].setdefault(key, {})[name] = (is_quantized, arr.copy())
+            self.tensors_dicts[-1].setdefault(node, {})[name] = (is_quantized, arr.copy())
 
     def pre_batch(self, m, b):
         self.tensors_dicts.append({})

test/adaptor/mxnet_adaptor/test_adaptor_mxnet.py

@@ -159,23 +159,23 @@ def test_mlp_model_quantization(self):
         fc_node_name2 = fc_op_name + '_1'
 
         insp = inspect_tensor(quantizer.model, quantizer.calib_dataloader,
-                              op_list=[(fc_node_name1, fc_op_name),
-                                       (fc_node_name2, fc_op_name)], iteration_list=[1, 3])
+                              op_list=[fc_node_name1, 
+                                       fc_node_name2 ], iteration_list=[1, 3])
         qinsp = inspect_tensor(qmodel, quantizer.calib_dataloader,
-                               op_list=[(fc_node_name1, fc_op_name),
-                                        (fc_node_name2, fc_op_name)], iteration_list=[1, 3])
+                               op_list=[fc_node_name1,
+                                        fc_node_name2], iteration_list=[1, 3])
 
         self.assertNotEqual(len(insp['activation']), 0)
         self.assertEqual(len(insp['activation']), len(qinsp['activation']))
 
         for tensors, qtensors in zip(insp['activation'], qinsp['activation']):
           for k in (set(tensors.keys()) & set(qtensors.keys())):
-            tensor, qtensor = tensors[k][k[0]], qtensors[k][k[0]]
+            tensor, qtensor = tensors[k][k], qtensors[k][k]
             self.assertEqual(tensor.shape, qtensor.shape)
 
         #test inspect with an empty iteration_list
         inspect_tensor(qmodel, quantizer.calib_dataloader,
-                       op_list=[(fc_node_name1, fc_op_name)],
+                       op_list=[fc_node_name1],
                        iteration_list=[])
 
         # test recovery for symbolic model