Consolidate Softmax range adjustment

microsoft · Mar 22, 2024 · d4cd04d · d4cd04d
1 parent ae3f3a3
commit d4cd04d
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/onnxruntime/python/tools/quantization/base_quantizer.py b/onnxruntime/python/tools/quantization/base_quantizer.py
@@ -452,7 +452,7 @@ def quantize_weight_per_channel_impl(
 
         return q_weight_name, zp_name, scale_name
 
-    def adjust_tensor_ranges(self, softmax_0_to_1=False):
+    def adjust_tensor_ranges(self):
         if self.tensors_range is None:
             return
 
@@ -471,6 +471,6 @@ def adjust_tensor_ranges(self, softmax_0_to_1=False):
                 if not isinstance(td, TensorData):
                     raise TypeError(f"Unexpected type {type(td)} for {node.output[0]!r}.")
                 self.tensors_range[node.input[0]] = td
-            # Optionally, adjust Softmax to range from 0.0 to 1.0
-            elif node.op_type == "Softmax" and softmax_0_to_1:
+            # Adjust Softmax to range from 0.0 to 1.0
+            elif node.op_type == "Softmax":
                 self.tensors_range[node.output[0]] = TensorData(lowest=np.float32(0.0), highest=np.float32(1.0))
diff --git a/onnxruntime/python/tools/quantization/onnx_quantizer.py b/onnxruntime/python/tools/quantization/onnx_quantizer.py
@@ -955,7 +955,7 @@ def calculate_quantization_params(self):
         if self.tensors_range is None:
             return None
 
-        self.adjust_tensor_ranges(softmax_0_to_1=False)
+        self.adjust_tensor_ranges()
 
         quantization_params = {}
         for tensor_name in self.tensors_range:

diff --git a/onnxruntime/python/tools/quantization/qdq_quantizer.py b/onnxruntime/python/tools/quantization/qdq_quantizer.py
@@ -1098,7 +1098,7 @@ def calc_graph_quant_params(self) -> dict[str, QDQTensorQuantParams]:
         if self.tensors_range is None:
             return {}
 
-        self.adjust_tensor_ranges(softmax_0_to_1=True)  # Ensure Softmax ranges from 0.0 to 1.0 for QDQ models.
+        self.adjust_tensor_ranges()
 
         quantization_params = {}
         for tensor_name in self.tensors_range: