fix diagnosis issue (#1328)

Author: mengniwang95

neural_compressor/adaptor/ox_utils/onnxrt_mid.py

@@ -183,7 +183,7 @@ def augment_graph(self, activation_only=False, weight_only=False):
 
         self.augmented_model = model
 
-    def get_intermediate_outputs(self):
+    def get_intermediate_outputs(self, calib_mode=None):
         '''
             Gather intermediate model outputs after running inference
             :return: dictionary mapping: {node output tensor names: node output tensor }
@@ -228,14 +228,20 @@ def get_intermediate_outputs(self):
                     break
                 if idx in self.iterations:
                     for output_idx, output in enumerate(session.run(None, ort_inputs)): 
-                        if output.size != 0:
+                        if calib_mode == 'naive' and output.size != 0:
                             output_dicts.setdefault(node_output_names[output_idx], \
                                 []).append([output.min(), output.max()])
+                        elif calib_mode == None:
+                            output_dicts.setdefault(node_output_names[output_idx], \
+                                []).append(output)
             else:
                 for output_idx, output in enumerate(session.run(None, ort_inputs)): 
-                    if output.size != 0:
+                    if calib_mode == 'naive' and output.size != 0:
                         output_dicts.setdefault(node_output_names[output_idx], \
                             []).append([output.min(), output.max()])
+                    elif calib_mode == None:
+                        output_dicts.setdefault(node_output_names[output_idx], \
+                            []).append(output)
 
         return list(output_dicts.keys()), output_dicts
 
@@ -344,7 +350,7 @@ def _map_calibration(self, node_output_names, output_dicts, calib_mode='naive'):
 
     def dump_minmax(self, calib_mode='naive'):
         self.augment_graph()
-        node_output_names, output_dicts = self.get_intermediate_outputs()
+        node_output_names, output_dicts = self.get_intermediate_outputs(calib_mode)
         return self._map_calibration(node_output_names, output_dicts,
                                      calib_mode=calib_mode)
 

test/adaptor/onnxrt_adaptor/test_onnxrt_augment.py

@@ -503,7 +503,7 @@ def test_quant_param_calculation(self):
         #test calculation of quantization params
         #TO_DO: check rmin/rmax
         quantization_params_dict = augment.dump_calibration()
-        node_output_names, output_dicts_list = augment.get_intermediate_outputs()
+        node_output_names, output_dicts_list = augment.get_intermediate_outputs('naive')
         dict_for_quantization = augment._map_calibration(node_output_names, output_dicts_list)
         #check the size of the quantization dictionary
         self.assertEqual(len(quantization_params_dict), 12)