mct-nightly 1.10.0.20231203.post417__py3-none-any.whl → 1.10.0.20231205.post412__py3-none-any.whl

{mct_nightly-1.10.0.20231203.post417.dist-info → mct_nightly-1.10.0.20231205.post412.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 1.10.0.20231203.post417
+Version: 1.10.0.20231205.post412
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct_nightly-1.10.0.20231203.post417.dist-info → mct_nightly-1.10.0.20231205.post412.dist-info}/RECORD

@@ -6,7 +6,7 @@ model_compression_toolkit/core/analyzer.py,sha256=dbsD61pakp_9JXNyAScLdtJvcXny9j
 model_compression_toolkit/core/exporter.py,sha256=U_-ea-zYHsnIt2ydameMLZ_gzDaCMI1dRa5IjA8RUuc,4233
 model_compression_toolkit/core/graph_prep_runner.py,sha256=SHhFl0vpC9YpRu40xkApFzmw_dT-nfIz1MDjmKcon8Q,9913
 model_compression_toolkit/core/quantization_prep_runner.py,sha256=npv55-QsJFR7bnbHj4tBMf13Y18Ns7QGa-UDSI6WJRE,6554
-model_compression_toolkit/core/runner.py,sha256=D_5OZ1PHlvytjCwu_7nTPkYHdmthN-ZHANpirqrH5A8,12069
+model_compression_toolkit/core/runner.py,sha256=Cb8_TWAOBz4SO1O48ehxqC9PpaR4KifbCs0nV724zMM,10454
 model_compression_toolkit/core/common/__init__.py,sha256=Wh127PbXcETZX_d1PQqZ71ETK3J9XO5A-HpadGUbj6o,1447
 model_compression_toolkit/core/common/base_substitutions.py,sha256=xDFSmVVs_iFSZfajytI0cuQaNRNcwHX3uqOoHgVUvxQ,1666
 model_compression_toolkit/core/common/data_loader.py,sha256=7YF5Mqz64Xb4rVwY3knrdIZ4JEHybXxiQqx0deR_c5k,4017
@@ -18,7 +18,7 @@ model_compression_toolkit/core/common/model_builder_mode.py,sha256=jll9-59OPaE3u
 model_compression_toolkit/core/common/model_collector.py,sha256=pNmJsU7QPCQ8-YUrzz__85YwF7Mk4Q27gozDSYCpzrg,5005
 model_compression_toolkit/core/common/model_validation.py,sha256=LaG8wd6aZl0OJgieE3SeiVDEPxtk8IHq9-3wSnmWhY4,1214
 model_compression_toolkit/core/common/node_prior_info.py,sha256=WXX_PrGVG9M9I_REG5ZzFBohwmV4yf356sZnrja_FLo,2832
-model_compression_toolkit/core/common/similarity_analyzer.py,sha256=QEwH9BZ7qKPt-x_1B_T_D014_B0m5z21Yz8c1Lmu-WM,8398
+model_compression_toolkit/core/common/similarity_analyzer.py,sha256=v7TBF_6iufFPTRQMT74Z2BuwSN4HfX7eFuqH4AF6JOM,8579
 model_compression_toolkit/core/common/user_info.py,sha256=BM98W0jRF_M6zYdCs6z7eKgEOFGu1DZjJja7KSSUGJQ,1631
 model_compression_toolkit/core/common/back2framework/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/core/common/back2framework/base_model_builder.py,sha256=V1oShKzbSkdcTvREn8VnQQBzvm-tTHkWMXqMkYozF2s,2023
@@ -132,7 +132,7 @@ model_compression_toolkit/core/common/substitutions/weights_activation_split.py,
 model_compression_toolkit/core/common/visualization/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/common/visualization/final_config_visualizer.py,sha256=6I10jKLesB-RQKaXA75Xgz2wPvylQUrnPtCcQZIynGo,6371
 model_compression_toolkit/core/common/visualization/nn_visualizer.py,sha256=6EjZj_KE1tICTQ0XSKIx5ivsRFpRktFywda7pW7YnNQ,5955
-model_compression_toolkit/core/common/visualization/tensorboard_writer.py,sha256=1Vbr5gKqH1fJl91iTNFlIjyEMh6jm88T4AIWalMrJFw,20099
+model_compression_toolkit/core/common/visualization/tensorboard_writer.py,sha256=954742gUTrrKmcVjcuBJaKR-EfMMsrWZ7PXd07unA6E,21939
 model_compression_toolkit/core/keras/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/keras/constants.py,sha256=oFYFagoFTOQTrs2RHVc93583EhOvcvbCYHleqsZdQ6s,3046
 model_compression_toolkit/core/keras/custom_layer_validation.py,sha256=f-b14wuiIgitBe7d0MmofYhDCTO3IhwJgwrh-Hq_t_U,1192
@@ -318,7 +318,7 @@ model_compression_toolkit/gptq/keras/gptq_keras_implementation.py,sha256=axBwnCS
 model_compression_toolkit/gptq/keras/gptq_loss.py,sha256=rbRkF15MYd6nq4G49kcjb_dPTa-XNq9cTkrb93mXawo,6241
 model_compression_toolkit/gptq/keras/gptq_training.py,sha256=75j56X2AcNI_5hInsLvXnWZOGMZIKQY2hStVKBaA_Bc,17705
 model_compression_toolkit/gptq/keras/graph_info.py,sha256=FIGqzJbG6GkdHenvdMu-tGTjp4j9BewdF_spmWCb4Mo,4627
-model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=2AJ3--JEg7nN4e8mKzMWqNPViIPVXb01C9QSoG2okdc,15118
+model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=DKh16lW24btKCrS6U-itaTHuTRX39e_nV5gn3hKUfxQ,15198
 model_compression_toolkit/gptq/keras/quantizer/__init__.py,sha256=-DK1CDXvlsnEbki4lukZLpl6Xrbo91_jcqxXlG5Eg6Q,963
 model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py,sha256=8NrJBftKFbMAF_jYaAbLP6GBwpCv3Ln1NKURaV75zko,4770
 model_compression_toolkit/gptq/keras/quantizer/quant_utils.py,sha256=Vt7Qb8i4JsE4sFtcjpfM4FTXTtfV1t6SwfoNH8a_Iaw,5055
@@ -335,7 +335,7 @@ model_compression_toolkit/gptq/pytorch/gptq_loss.py,sha256=kDuWw-6zh17wZpYWh4Xa9
 model_compression_toolkit/gptq/pytorch/gptq_pytorch_implementation.py,sha256=tECPTavxn8EEwgLaP2zvxdJH6Vg9jC0YOIMJ7857Sdc,1268
 model_compression_toolkit/gptq/pytorch/gptq_training.py,sha256=RfkJNQbeqGZQvlmV0dZO7YJ894Gx2asLnnIHFdWNEZ0,15078
 model_compression_toolkit/gptq/pytorch/graph_info.py,sha256=-0GDC2cr-XXS7cTFTnDflJivGN7VaPnzVPsxCE-vZNU,3955
-model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=P2vEYIQvLKxOFX4bCUEPheq63QCkk6G20Lga8AaFJgo,13153
+model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=q1CzOc0LY1i8Ux4uH2PFGlGhTAJ0t-PPkIJURzbMu9w,13233
 model_compression_toolkit/gptq/pytorch/quantizer/__init__.py,sha256=ZHNHo1yzye44m9_ht4UUZfTpK01RiVR3Tr74-vtnOGI,968
 model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py,sha256=Zb-P0yRyZHHBlDvUBdRwxDpdduEJyJp6OT9pfKFF5ks,4171
 model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py,sha256=OocYYRqvl7rZ37QT0hTzfJnWGiNCPskg7cziTlR7TRk,3893
@@ -348,19 +348,19 @@ model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/uniform_soft_quan
 model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/symmetric_ste.py,sha256=whVx94NCmYzobvZtPNr4qZTCN-CV8jfs4des_mkK3F8,8770
 model_compression_toolkit/legacy/__init__.py,sha256=lNJ29DYxaLUPDstRDA1PGI5r9Fulq_hvrZMlhst1Z5g,697
-model_compression_toolkit/legacy/keras_quantization_facade.py,sha256=r9iHnpdROkLN1aiCOVOCJlLC9Z-ZNIT8wOFyhn0ACvY,18365
-model_compression_toolkit/legacy/pytorch_quantization_facade.py,sha256=LwdlzAVs2xkdPUuizWrE7sQkHZmY2G5v31uahl4gcSw,17795
+model_compression_toolkit/legacy/keras_quantization_facade.py,sha256=2pNJoc1mKMbikBS_uebLgFAbTqfA0y9ofDUNCVogSKI,18444
+model_compression_toolkit/legacy/pytorch_quantization_facade.py,sha256=p-ZGKdGeRIJsR5XmFYgjs3VN49NrwHumNtTY2OSDW-4,17874
 model_compression_toolkit/ptq/__init__.py,sha256=50QBTXOKdj9XLjXtrvf0mhC9FlW6TOi9-pjl96RLR14,930
 model_compression_toolkit/ptq/runner.py,sha256=_c1dSjlPPpsx59Vbg1buhG9bZq__OORz1VlPkwjJzoc,2552
 model_compression_toolkit/ptq/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/ptq/keras/quantization_facade.py,sha256=OKycZVFXDFrxOFY1IjLvIhG_kBXBDtrRsz54n_FYMMo,9889
+model_compression_toolkit/ptq/keras/quantization_facade.py,sha256=39-bHTOZ7OENxvbC_HKelwjSO3e8BccGexgwWPClIDk,9969
 model_compression_toolkit/ptq/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=nsfvvLbgERes9qqNSLi-X5EgCIhzpLBDqK_dCAKAbGk,8670
+model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=zsewr833uH91wpluQ4Z22B3HX5ohepdJZPTqNpAosZk,8750
 model_compression_toolkit/qat/__init__.py,sha256=BYKgH1NwB9fqF1TszULQ5tDfLI-GqgZV5sao-lDN9EM,1091
 model_compression_toolkit/qat/common/__init__.py,sha256=6tLZ4R4pYP6QVztLVQC_jik2nES3l4uhML0qUxZrezk,829
 model_compression_toolkit/qat/common/qat_config.py,sha256=kbSxFL6_u28furq5mW_75STWDmyX4clPt-seJAnX3IQ,3445
 model_compression_toolkit/qat/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/qat/keras/quantization_facade.py,sha256=bH9pw2RtFD_91pfA5LPOuZ_ot04uX8R1au_n3Vr_7-0,16173
+model_compression_toolkit/qat/keras/quantization_facade.py,sha256=hB_VrhSqwOjGnOT8BYpXkh52EMzo7I-62-IEGohg_74,16253
 model_compression_toolkit/qat/keras/quantizer/__init__.py,sha256=zmYyCa25_KLCSUCGUDRslh3RCIjcRMxc_oXa54Aui-4,996
 model_compression_toolkit/qat/keras/quantizer/base_keras_qat_quantizer.py,sha256=gPuIgQb8OafvC3SuA8jNsGoy8S8eTsDCEKuh36WDNss,2104
 model_compression_toolkit/qat/keras/quantizer/quant_utils.py,sha256=cBULOgWUodcBO1lHevZggdTevuDYI6tQceV86U2x6DA,2543
@@ -372,7 +372,7 @@ model_compression_toolkit/qat/keras/quantizer/ste_rounding/__init__.py,sha256=cc
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/symmetric_ste.py,sha256=I4KlaGv17k71IyjuSG9M0OlXlD5P0pfvKa6oCyRQ5FE,13517
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/uniform_ste.py,sha256=EED6LfqhX_OhDRJ9e4GwbpgNC9vq7hoXyJS2VPvG2qc,10789
 model_compression_toolkit/qat/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=LGIk_nbcQEOhSN8fe56u6QGGUtGOVhF7tlCwW2G-Ig4,12549
+model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=Kv2vUKY2cSt6kYB7J2B1SC_PGzSFqecpxDcKWGeRzuQ,12629
 model_compression_toolkit/qat/pytorch/quantizer/__init__.py,sha256=xYa4C8pr9cG1f3mQQcBXO_u3IdJN-zl7leZxuXDs86w,1003
 model_compression_toolkit/qat/pytorch/quantizer/base_pytorch_qat_quantizer.py,sha256=FnhuFCuQoSf78FM1z1UZgXXd3k-mKSM7i9dYOuJUmeA,2213
 model_compression_toolkit/qat/pytorch/quantizer/quantization_builder.py,sha256=GOYRDXvQSGe_iUFVmvDy5BqC952hu_-rQO06n8QCyw0,5491
@@ -448,8 +448,8 @@ model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=SbvRlIdE32PEBsINt1bhSqvrKL_zbM9V-aeSkOn-sw4,3083
-mct_nightly-1.10.0.20231203.post417.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-1.10.0.20231203.post417.dist-info/METADATA,sha256=awXN4yl5z3uUvO-0tUbhzQDIKW68lc9Lbjji5pJ0zlM,16232
-mct_nightly-1.10.0.20231203.post417.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-mct_nightly-1.10.0.20231203.post417.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-1.10.0.20231203.post417.dist-info/RECORD,,
+mct_nightly-1.10.0.20231205.post412.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-1.10.0.20231205.post412.dist-info/METADATA,sha256=TJLsSdD5_ptD9k1nxSMBU83HMkoDdwh1OGP9_dMrhbM,16232
+mct_nightly-1.10.0.20231205.post412.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+mct_nightly-1.10.0.20231205.post412.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-1.10.0.20231205.post412.dist-info/RECORD,,

model_compression_toolkit/core/common/similarity_analyzer.py

@@ -241,4 +241,7 @@ def compute_kl_divergence(float_tensor: np.ndarray, fxp_tensor: np.ndarray, batc
     non_zero_fxp_tensor = fxp_flat.copy()
     non_zero_fxp_tensor[non_zero_fxp_tensor == 0] = EPS
 
-    return np.mean(np.sum(np.where(float_flat != 0, float_flat * np.log(float_flat / non_zero_fxp_tensor), 0), axis=-1))
+    prob_distance = np.where(float_flat != 0, float_flat * np.log(float_flat / non_zero_fxp_tensor), 0)
+    # The sum is part of the KL-Divergance function.
+    # The mean is to aggregate the distance between each output probability vectors.
+    return np.mean(np.sum(prob_distance, axis=-1), axis=-1)

model_compression_toolkit/core/common/visualization/tensorboard_writer.py

@@ -16,6 +16,7 @@
 from copy import deepcopy
 
 import io
+import os
 import numpy as np
 from PIL import Image
 from matplotlib.figure import Figure
@@ -34,6 +35,9 @@ from networkx import topological_sort
 from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.collectors.statistics_collector import BaseStatsCollector
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.core.common.visualization.final_config_visualizer import \
+    WeightsFinalBitwidthConfigVisualizer, ActivationFinalBitwidthConfigVisualizer
 
 DEVICE_STEP_STATS = "/device:CPU:0"
 
@@ -486,3 +490,45 @@ class TensorboardWriter(object):
         er = self.__get_event_writer_by_tag_name(main_tag_name)
         er.add_event(event)
         er.flush()
+
+
+def init_tensorboard_writer(fw_info: FrameworkInfo) -> TensorboardWriter:
+    """
+    Create a TensorBoardWriter object initialized with the logger dir path if it was set,
+    or None otherwise.
+
+    Args:
+        fw_info: FrameworkInfo object.
+
+    Returns:
+        A TensorBoardWriter object.
+    """
+    tb_w = None
+    if Logger.LOG_PATH is not None:
+        tb_log_dir = os.path.join(os.getcwd(), Logger.LOG_PATH, 'tensorboard_logs')
+        Logger.info(f'To use Tensorboard, please run: tensorboard --logdir {tb_log_dir}')
+        tb_w = TensorboardWriter(tb_log_dir, fw_info)
+    return tb_w
+
+
+def finalize_bitwidth_in_tb(tb_w: TensorboardWriter,
+                            weights_conf_nodes_bitwidth: List,
+                            activation_conf_nodes_bitwidth: List):
+    """
+    Set the final bit-width configuration of the quantized model in the provided TensorBoard object.
+
+    Args:
+        tb_w: A TensorBoard object.
+        weights_conf_nodes_bitwidth: Final weights bit-width configuration.
+        activation_conf_nodes_bitwidth: Final activation bit-width configuration.
+
+    """
+
+    if len(weights_conf_nodes_bitwidth) > 0:
+        visual = WeightsFinalBitwidthConfigVisualizer(weights_conf_nodes_bitwidth)
+        figure = visual.plot_config_bitwidth()
+        tb_w.add_figure(figure, f'Weights final bit-width config')
+    if len(activation_conf_nodes_bitwidth) > 0:
+        visual = ActivationFinalBitwidthConfigVisualizer(activation_conf_nodes_bitwidth)
+        figure = visual.plot_config_bitwidth()
+        tb_w.add_figure(figure, f'Activation final bit-width config')

model_compression_toolkit/core/runner.py

@@ -14,11 +14,9 @@
 # ==============================================================================
 
 
-import os
 from typing import Callable, Tuple, Any, List, Dict
 
 import numpy as np
-from tqdm import tqdm
 
 from model_compression_toolkit.core.common import FrameworkInfo
 from model_compression_toolkit.core.common.hessian.hessian_info_service import HessianInfoService
@@ -33,20 +31,14 @@ from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi_aggrega
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi_functions_mapping import kpi_functions_mapping
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi_methods import MpKpiMetric
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_search_facade import search_bit_width
-from model_compression_toolkit.core.common.model_collector import ModelCollector
 from model_compression_toolkit.core.common.network_editors.edit_network import edit_network_graph
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
-from model_compression_toolkit.core.common.quantization.quantization_analyzer import analyzer_graph
-from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_computation import \
-    calculate_quantization_params
-from model_compression_toolkit.core.common.statistics_correction.statistics_correction import \
-    statistics_correction_runner
-from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.visualization.final_config_visualizer import \
     WeightsFinalBitwidthConfigVisualizer, \
     ActivationFinalBitwidthConfigVisualizer
-from model_compression_toolkit.core.common.visualization.tensorboard_writer import TensorboardWriter
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import TensorboardWriter, \
+    finalize_bitwidth_in_tb
 
 
 def core_runner(in_model: Any,
@@ -150,37 +142,11 @@ def core_runner(in_model: Any,
             f'Final activation bit-width configuration: {[node_b[1] for node_b in activation_conf_nodes_bitwidth]}')
 
         if tb_w is not None:
-            if len(weights_conf_nodes_bitwidth) > 0:
-                visual = WeightsFinalBitwidthConfigVisualizer(weights_conf_nodes_bitwidth)
-                figure = visual.plot_config_bitwidth()
-                tb_w.add_figure(figure, f'Weights final bit-width config')
-            if len(activation_conf_nodes_bitwidth) > 0:
-                visual = ActivationFinalBitwidthConfigVisualizer(activation_conf_nodes_bitwidth)
-                figure = visual.plot_config_bitwidth()
-                tb_w.add_figure(figure, f'Activation final bit-width config')
+            finalize_bitwidth_in_tb(tb_w, weights_conf_nodes_bitwidth, activation_conf_nodes_bitwidth)
 
     return tg, bit_widths_config, hessian_info_service
 
 
-def _init_tensorboard_writer(fw_info: FrameworkInfo) -> TensorboardWriter:
-    """
-    Create a TensorBoardWriter object initialized with the logger dir path if it was set,
-    or None otherwise.
-
-    Args:
-        fw_info: FrameworkInfo object.
-
-    Returns:
-        A TensorBoardWriter object.
-    """
-    tb_w = None
-    if Logger.LOG_PATH is not None:
-        tb_log_dir = os.path.join(os.getcwd(), Logger.LOG_PATH, 'tensorboard_logs')
-        Logger.info(f'To use Tensorboard, please run: tensorboard --logdir {tb_log_dir}')
-        tb_w = TensorboardWriter(tb_log_dir, fw_info)
-    return tb_w
-
-
 def _set_final_kpi(graph: Graph,
                    final_bit_widths_config: List[int],
                    kpi_functions_dict: Dict[KPITarget, Tuple[MpKpiMetric, MpKpiAggregation]],

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -16,6 +16,7 @@
 from typing import Callable, Tuple
 from packaging import version
 
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW, FOUND_TF
 from model_compression_toolkit.core.common.user_info import UserInformation
@@ -24,7 +25,7 @@ from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfigV2
 from model_compression_toolkit.core import CoreConfig
-from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
+from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.core.exporter import export_model
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
@@ -202,7 +203,7 @@ if FOUND_TF:
             Logger.info("Using experimental mixed-precision quantization. "
                                "If you encounter an issue please file a bug.")
 
-        tb_w = _init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = GPTQKerasImplemantation()
 

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -15,12 +15,13 @@
 from typing import Callable
 from model_compression_toolkit.core import common
 from model_compression_toolkit.constants import FOUND_TORCH
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import PYTORCH
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
 from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
-from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
+from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.gptq.keras.quantization_facade import GPTQ_MOMENTUM
 from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.core.exporter import export_model
@@ -161,7 +162,7 @@ if FOUND_TORCH:
             Logger.info("Using experimental mixed-precision quantization. "
                                "If you encounter an issue please file a bug.")
 
-        tb_w = _init_tensorboard_writer(DEFAULT_PYTORCH_INFO)
+        tb_w = init_tensorboard_writer(DEFAULT_PYTORCH_INFO)
 
         fw_impl = GPTQPytorchImplemantation()
 

model_compression_toolkit/legacy/keras_quantization_facade.py

@@ -15,6 +15,7 @@
 
 from typing import Callable, List, Tuple
 
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW
 from model_compression_toolkit.core.common.user_info import UserInformation
@@ -28,7 +29,7 @@ from model_compression_toolkit.core.common.quantization.quantization_config impo
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.quantization.debug_config import DebugConfig
 from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG
-from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
+from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.ptq.runner import ptq_runner
 from model_compression_toolkit.core.exporter import export_model
@@ -114,7 +115,7 @@ if FOUND_TF:
                                                           network_editor=network_editor)
                                  )
 
-        tb_w = _init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = KerasImplementation()
 
@@ -249,7 +250,7 @@ if FOUND_TF:
                                                           network_editor=network_editor)
                                  )
 
-        tb_w = _init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = KerasImplementation()
 

model_compression_toolkit/legacy/pytorch_quantization_facade.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 from typing import Callable, List, Tuple
 
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import PYTORCH
 from model_compression_toolkit.core.common.user_info import UserInformation
@@ -28,7 +29,7 @@ from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quant
     MixedPrecisionQuantizationConfig, DEFAULT_MIXEDPRECISION_CONFIG
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
 from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG
-from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
+from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.ptq.runner import ptq_runner
 from model_compression_toolkit.core.exporter import export_model
@@ -106,7 +107,7 @@ if FOUND_TORCH:
                                  debug_config=DebugConfig(analyze_similarity=analyze_similarity,
                                                           network_editor=network_editor))
 
-        tb_w = _init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = PytorchImplementation()
 
@@ -235,7 +236,7 @@ if FOUND_TORCH:
                                  debug_config=DebugConfig(analyze_similarity=analyze_similarity,
                                                           network_editor=network_editor))
 
-        tb_w = _init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = PytorchImplementation()
 

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -17,6 +17,7 @@ from typing import Callable
 
 from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW, FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
@@ -24,7 +25,7 @@ from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quant
     MixedPrecisionQuantizationConfigV2
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
 from model_compression_toolkit.core.exporter import export_model
-from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
+from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.ptq.runner import ptq_runner
 
 if FOUND_TF:
@@ -130,7 +131,7 @@ if FOUND_TF:
             Logger.info("Using experimental mixed-precision quantization. "
                                "If you encounter an issue please file a bug.")
 
-        tb_w = _init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = KerasImplementation()
 

model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -15,6 +15,7 @@
 from typing import Callable
 
 from model_compression_toolkit.core import common
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import PYTORCH, FOUND_TORCH
 from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
@@ -22,7 +23,7 @@ from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import
 from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfigV2
-from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
+from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.ptq.runner import ptq_runner
 from model_compression_toolkit.core.exporter import export_model
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
@@ -102,7 +103,7 @@ if FOUND_TORCH:
             Logger.info("Using experimental mixed-precision quantization. "
                                "If you encounter an issue please file a bug.")
 
-        tb_w = _init_tensorboard_writer(DEFAULT_PYTORCH_INFO)
+        tb_w = init_tensorboard_writer(DEFAULT_PYTORCH_INFO)
 
         fw_impl = PytorchImplementation()
 

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -17,6 +17,7 @@ from typing import Callable
 from functools import partial
 
 from model_compression_toolkit.core import CoreConfig
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
@@ -25,7 +26,7 @@ from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quant
 from mct_quantizers import KerasActivationQuantizationHolder
 from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
-from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
+from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.ptq.runner import ptq_runner
 
 if FOUND_TF:
@@ -177,7 +178,7 @@ if FOUND_TF:
             Logger.info("Using experimental mixed-precision quantization. "
                                "If you encounter an issue please file a bug.")
 
-        tb_w = _init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = KerasImplementation()
 

model_compression_toolkit/qat/pytorch/quantization_facade.py

@@ -20,6 +20,7 @@ from model_compression_toolkit.constants import FOUND_TORCH, PYTORCH
 
 from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core import common
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
@@ -27,7 +28,7 @@ from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quant
     MixedPrecisionQuantizationConfigV2
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import \
     TargetPlatformCapabilities
-from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
+from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.ptq.runner import ptq_runner
 
 if FOUND_TORCH:
@@ -145,7 +146,7 @@ if FOUND_TORCH:
             Logger.info("Using experimental mixed-precision quantization. "
                         "If you encounter an issue please file a bug.")
 
-        tb_w = _init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = PytorchImplementation()