Skip to content

Commit

Permalink
Modified access of internal utilities
Browse files Browse the repository at this point in the history 
Signed-off-by: Raj Gite <quic_rgite@quicinc.com>
  • Loading branch information
@quic-rgite @quic-akhobare
quic-rgite authored and quic-akhobare committed Jan 30, 2024
1 parent 5185784 commit 37ff059
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 27 deletions.
46 changes: 27 additions & 19 deletions TrainingExtensions/onnx/src/python/aimet_onnx/quant_analyzer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -130,13 +130,13 @@ def analyze(self,
self.check_model_sensitivity_to_quantization(sim)

# Perform per layer analysis by enabling its quantizers (OPTION-1).
self._perform_per_layer_analysis_by_enabling_quantizers(sim, results_dir)
self.perform_per_layer_analysis_by_enabling_quantizers(sim, results_dir)

# Perform per layer analysis by disabling its quantizers (OPTION-2).
self._perform_per_layer_analysis_by_disabling_quantizers(sim, results_dir)
self.perform_per_layer_analysis_by_disabling_quantizers(sim, results_dir)

def create_quantsim_and_encodings(self, quant_scheme: QuantScheme, default_param_bw: int,
default_activation_bw: int, config_file: str) \
default_activation_bw: int, config_file: str) \
-> QuantizationSimModel:
"""
Create Quantsim and compute encodings.
Expand Down Expand Up @@ -262,10 +262,10 @@ def _enable_disable_quantizers(quantizers: List[QcQuantizeOp], enabled: bool):
for quantizer in quantizers:
quantizer.enabled = enabled

def _perform_per_layer_analysis_by_enabling_quantizers(self,
sim: QuantizationSimModel,
results_dir: str,
) -> Dict:
def perform_per_layer_analysis_by_enabling_quantizers(self,
sim: QuantizationSimModel,
results_dir: str,
) -> Dict:
"""
NOTE: Option 1
Expand Down Expand Up @@ -299,10 +299,10 @@ def _perform_per_layer_analysis_by_enabling_quantizers(self,
_logger.info("Exported per-layer quant analysis (enabled) plot.")
return layer_wise_eval_score_dict

def _perform_per_layer_analysis_by_disabling_quantizers(self,
sim: QuantizationSimModel,
results_dir: str,
) -> Dict:
def perform_per_layer_analysis_by_disabling_quantizers(self,
sim: QuantizationSimModel,
results_dir: str,
) -> Dict:
"""
NOTE: Option 2
Expand Down Expand Up @@ -343,8 +343,8 @@ def _perform_per_layer_analysis(self,
enabled_after: bool,
) -> Dict:
"""
Helper function for perform_per_layer_analysis_by_enabling_quant_wrappers() and
perform_per_layer_analysis_by_disabling_quant_wrappers()
Helper function for perform_per_layer_analysis_by_enabling_quantizers() and
perform_per_layer_analysis_by_disabling_quantizers()
:param sim: Quantsim model.
:param disable_all_quantizers: Flag to disable all the quantizers before per-layer analysis.
Expand Down Expand Up @@ -377,28 +377,36 @@ def _perform_per_layer_analysis(self,
return eval_score_dict

@staticmethod
def _get_enabled_quantizers(sim: QuantizationSimModel):
enabled_quant_wrappers = defaultdict(list)
def _get_enabled_quantizers(sim: QuantizationSimModel) -> Dict:
"""
This function creates mapping between ops and their enabled quantizers.
:param sim: Quantsim object
:return: Layerwise enabled quantizers
"""
enabled_quantizers = defaultdict(list)
cg_ops = sim.connected_graph.ordered_ops
for op in cg_ops:
# Get param quantizers
for param in op.parameters:
if param in sim.qc_quantize_op_dict and sim.qc_quantize_op_dict[param].enabled:
enabled_quant_wrappers[op.name_op].append(sim.qc_quantize_op_dict[param])
enabled_quantizers[op.name_op].append(sim.qc_quantize_op_dict[param])
# Get output activation quantizers
if op.output_ops and op.output_ops[0].type == 'branch':
# op having multiple outputs
cg_product = op.output_ops[0].output
else:
# op having single output
cg_product = op.output
for output_name in set(cg_product.tensor_dict.values()):
if output_name in sim.qc_quantize_op_dict and sim.qc_quantize_op_dict[output_name].enabled:
enabled_quant_wrappers[op.name_op].append(sim.qc_quantize_op_dict[output_name])
enabled_quantizers[op.name_op].append(sim.qc_quantize_op_dict[output_name])
# Get input activation quantizers if starting op
if op in sim.connected_graph.starting_ops:
cg_products = [cg_product for cg_product in op.inputs if cg_product.is_model_input]
for cg_product in cg_products:
assert len(cg_product.tensor_dict) == 1
input_name = list(cg_product.tensor_dict.values())[0]
if input_name in sim.qc_quantize_op_dict and sim.qc_quantize_op_dict[input_name].enabled:
enabled_quant_wrappers[op.name_op].append(sim.qc_quantize_op_dict[input_name])
return enabled_quant_wrappers
enabled_quantizers[op.name_op].append(sim.qc_quantize_op_dict[input_name])
return enabled_quantizers
16 changes: 8 additions & 8 deletions TrainingExtensions/onnx/test/python/test_quant_analyzer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -167,17 +167,17 @@ def test_perform_per_layer_analysis_by_enabling_quantizers(self):
quant_analyzer = QuantAnalyzer(model, dummy_input_dict, forward_pass_callback, eval_callback)
try:
layer_wise_eval_score_dict = \
quant_analyzer._perform_per_layer_analysis_by_enabling_quantizers(sim, results_dir="./tmp/")
quant_analyzer.perform_per_layer_analysis_by_enabling_quantizers(sim, results_dir="./tmp/")
print(layer_wise_eval_score_dict)
assert type(layer_wise_eval_score_dict) == dict
assert len(layer_wise_eval_score_dict) == 10

# test whether layer_wise_eval_score_dict consists of correct keys (op names).
# Test whether layer_wise_eval_score_dict consists of correct keys (op names).
for op_name in layer_wise_eval_score_dict.keys():
assert op_name in layer_names

# Check if it is exported to correct html file.
assert os.path.isfile("./tmp/per_layer_quant_enabled.html")
# Check if it is exported to correct html file.
assert os.path.isfile("./tmp/per_layer_quant_enabled.html")
finally:
if os.path.isdir("./tmp/"):
shutil.rmtree("./tmp/")
Expand All @@ -200,17 +200,17 @@ def test_perform_per_layer_analysis_by_disabling_quantizers(self):
quant_analyzer = QuantAnalyzer(model, dummy_input_dict, forward_pass_callback, eval_callback)
try:
layer_wise_eval_score_dict = \
quant_analyzer._perform_per_layer_analysis_by_disabling_quantizers(sim, results_dir="./tmp/")
quant_analyzer.perform_per_layer_analysis_by_disabling_quantizers(sim, results_dir="./tmp/")
print(layer_wise_eval_score_dict)
assert type(layer_wise_eval_score_dict) == dict
assert len(layer_wise_eval_score_dict) == 10

# test whether layer_wise_eval_score_dict consists of correct keys (op names).
# Test whether layer_wise_eval_score_dict consists of correct keys (op names).
for op_name in layer_wise_eval_score_dict.keys():
assert op_name in layer_names

# Check if it is exported to correct html file.
assert os.path.isfile("./tmp/per_layer_quant_disabled.html")
# Check if it is exported to correct html file.
assert os.path.isfile("./tmp/per_layer_quant_disabled.html")
finally:
if os.path.isdir("./tmp/"):
shutil.rmtree("./tmp/")

0 comments on commit 37ff059

Please sign in to comment.