Add progress bar for woq and sq (#1467)

mengniwang95 · yuwenzho · web-flow · commit 4d26e3aa94a2 · 2023-12-18T10:10:39.000+08:00
Co-authored-by: yuwenzho &lt;yuwen.zhou@intel.com&gt;
diff --git a/neural_compressor/adaptor/ox_utils/calibration.py b/neural_compressor/adaptor/ox_utils/calibration.py
@@ -757,6 +757,7 @@ def calib_smooth(self, percentile, op_types, q_config):
             max_vals_per_channel: max values per channel of input tensors
             shape_infos: The shape information of input tensors
         """
+        logger.info("Start smooth model calibration.")
         # add the input tensors of {op_types} to outputs of the model
         tensors_to_node = self._get_input_tensor_of_ops(op_types)
         self.model_wrapper.add_tensors_to_outputs(tensors_to_node.keys())
diff --git a/neural_compressor/adaptor/ox_utils/smooth_quant.py b/neural_compressor/adaptor/ox_utils/smooth_quant.py
@@ -25,7 +25,13 @@
 from onnx import helper, numpy_helper
 from onnx import onnx_pb as onnx_proto
 
-from neural_compressor.adaptor.ox_utils.util import _get_qrange_for_qType, is_B_transposed, quantize_data, to_numpy
+from neural_compressor.adaptor.ox_utils.util import (
+    _get_qrange_for_qType,
+    is_B_transposed,
+    quantize_data,
+    simple_progress_bar,
+    to_numpy,
+)
 from neural_compressor.model.model import BaseModel
 from neural_compressor.model.onnx_model import ONNXModel
 
@@ -562,6 +568,7 @@ def _get_smooth_scales(self, alpha, target_list=[]):
         Returns:
             the smooth scales for weights, currently one input tensor only have one scale
         """
+        logger.info("Start smooth scales collection.")
         scales = {}
         for tensor, nodes in self.tensors_to_node.items():
             # if scales_per_op the key of scales is the node name, otherwise the activation of node
@@ -673,7 +680,8 @@ def _adjust_weights(self, scales):
         Args:
             scales (dict): The input scales
         """
-        for tensor_name, nodes in self.tensors_to_node.items():
+        for idx, (tensor_name, nodes) in enumerate(self.tensors_to_node.items()):
+            simple_progress_bar(len(self.tensors_to_node), idx + 1)
             for node_info in nodes:
                 key = node_info[0] if self.scales_per_op else tensor_name
                 if key not in scales:
diff --git a/neural_compressor/adaptor/ox_utils/util.py b/neural_compressor/adaptor/ox_utils/util.py
@@ -93,6 +93,16 @@
 MAXIMUM_PROTOBUF = 2147483648
 
 
+def simple_progress_bar(total, i):
+    """Progress bar for cases where tqdm can't be used."""
+    progress = i / total
+    bar_length = 20
+    bar = "#" * int(bar_length * progress)
+    spaces = " " * (bar_length - len(bar))
+    percentage = progress * 100
+    print(f"\rProgress: [{bar}{spaces}] {percentage:.2f}%", end="")
+
+
 def dtype_to_name(dtype_mapping, dtype):
     """Map data type and its string representation."""
     return list(dtype_mapping.keys())[list(dtype_mapping.values()).index(dtype)]
diff --git a/neural_compressor/adaptor/ox_utils/weight_only.py b/neural_compressor/adaptor/ox_utils/weight_only.py
@@ -29,6 +29,7 @@
 from onnx import onnx_pb as onnx_proto
 from packaging.version import Version
 
+from neural_compressor.adaptor.ox_utils.util import simple_progress_bar
 from neural_compressor.model.model import BaseModel
 from neural_compressor.model.onnx_model import ONNXModel
 from neural_compressor.utils.utility import LazyImport
@@ -320,7 +321,12 @@ def rtn_quantize(
     base_dir = os.path.dirname(model.model_path) if model.model_path is not None else ""
     new_nodes = []
     remove_nodes = []
+    total_num = len([i for i in model.nodes() if i.op_type in ["MatMul"]])
+    curr_id = 0
     for node in model.nodes():
+        if node.op_type in ["MatMul"]:
+            curr_id += 1
+            simple_progress_bar(total_num, curr_id)
         if (
             node.op_type in ["MatMul"]
             and model.get_initializer(node.input[1]) is not None
@@ -1025,8 +1031,8 @@ def gptq_quantize(
 
     new_nodes = []
     remove_nodes = []
-
-    for input_name in output_names:
+    for idx, input_name in enumerate(output_names):
+        simple_progress_bar(len(output_names), idx + 1)
         node_list = []
         weights = []
 
