stm32ai_main.py

Go to the documentation of this file.

 
 
import os
import sys
from hydra.core.hydra_config import HydraConfig
import hydra
import warnings
warnings.filterwarnings("ignore")
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 
import tensorflow as tf
from omegaconf import DictConfig
import mlflow
import argparse
import logging
from clearml import Task
from clearml.backend_config.defs import get_active_config_file
 
 
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(SCRIPT_DIR))
 
from common.utils import mlflow_ini, set_gpu_memory_limit, get_random_seed, display_figures, log_to_file
from common.benchmarking import benchmark, cloud_connect
from common.evaluation import gen_load_val
from common.prediction import gen_load_val_predict
from src.preprocessing import preprocess
from src.utils import get_config
from src.training import train
from src.evaluation import evaluate
from src.quantization import quantize
from src.prediction import predict
from deployment import deploy, deploy_mpu
 
from typing import Optional
 
 
def chain_qd(cfg: DictConfig = None, float_model_path: str = None,
             train_ds: tf.data.Dataset = None,
             quantization_ds: tf.data.Dataset = None) -> None:
    """
    @brief Executes the Quantization → Deployment pipeline (chain_qd).
 
    @details
    This chain is used when a float model is already trained and only needs to be
    quantized and then deployed onto the STM32N6 board.
 
    Quantization strategy (in order of priority):
      1. Use the dedicated quantization dataset if provided.
      2. Fall back to the training dataset if no quantization dataset is available.
      3. Use fake (synthetic) data if neither dataset is provided — accuracy will be degraded.
 
    After quantization, the model is deployed:
      - On MPU targets via deploy_mpu().
      - On MCU targets (e.g., STM32N6570-DK) via deploy().
 
    @param cfg           Hydra configuration dictionary loaded from user_config.yaml.
    @param float_model_path  Path to the float32 model file (.tflite, .h5, or .onnx).
    @param train_ds      TensorFlow dataset used as fallback for quantization calibration.
    @param quantization_ds   Dedicated TensorFlow dataset for INT8 quantization calibration.
 
    @return None
    """
    # Optionally connect to STM32Cube.AI Developer Cloud for remote benchmarking/deployment
    credentials = None
    if cfg.tools.stm32ai.on_cloud:
        _, _, credentials = cloud_connect(stm32ai_version=cfg.tools.stm32ai.version)
 
    # Select the calibration dataset for quantization
    if quantization_ds:
        print('[INFO] : Using the quantization dataset to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=quantization_ds,
                                        float_model_path=float_model_path)
    elif train_ds:
        print('[INFO] : Quantization dataset is not provided! Using the training set to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=train_ds,
                                        float_model_path=float_model_path)
    else:
        print('[INFO] : Neither quantization dataset nor training set are provided! '
              'Using fake data to quantize the model. The model performance will not be accurate.')
        quantized_model_path = quantize(cfg=cfg, fake=True)
    print('[INFO] : Quantization complete.')
 
    # Deploy to MPU or MCU depending on hardware_type in the config
    if cfg.hardware_type == "MPU":
        deploy_mpu(cfg=cfg, model_path_to_deploy=quantized_model_path, credentials=credentials)
    else:
        deploy(cfg=cfg, model_path_to_deploy=quantized_model_path, credentials=credentials)
 
    print('[INFO] : Deployment complete.')
    # Remind the user to toggle boot switches on the STM32N6570-DK after flashing
    if cfg.deployment.hardware_setup.board == "STM32N6570-DK":
        print('[INFO] : On STM32N6570-DK, please toggle the boot switches to the left and power cycle the board.')
 
 
def chain_eqeb(cfg: DictConfig = None, float_model_path: str = None,
               train_ds: tf.data.Dataset = None,
               valid_ds: tf.data.Dataset = None,
               quantization_ds: tf.data.Dataset = None,
               test_ds: tf.data.Dataset = None) -> None:
    """
    @brief Executes the Evaluation → Quantization → Evaluation → Benchmarking pipeline (chain_eqeb).
 
    @details
    This chain is used to fully characterize both the float and quantized versions of a model:
      1. Evaluate the float model to establish a baseline accuracy.
      2. Quantize to INT8 using the provided calibration dataset.
      3. Evaluate the quantized model to measure accuracy degradation.
      4. Benchmark the quantized model on the target STM32 board to measure real-world latency.
 
    @param cfg               Hydra configuration dictionary.
    @param float_model_path  Path to the float32 model.
    @param train_ds          Training dataset (used as fallback for quantization calibration).
    @param valid_ds          Validation dataset for evaluation.
    @param quantization_ds   Dedicated calibration dataset for INT8 quantization.
    @param test_ds           Test dataset (takes priority over valid_ds for evaluation).
 
    @return None
    """
    credentials = None
    if cfg.tools.stm32ai.on_cloud:
        _, _, credentials = cloud_connect(stm32ai_version=cfg.tools.stm32ai.version)
 
    # Step 1: Evaluate the float model
    if test_ds:
        evaluate(cfg=cfg, eval_ds=test_ds, model_path_to_evaluate=float_model_path, name_ds="test_set")
    else:
        evaluate(cfg=cfg, eval_ds=valid_ds, model_path_to_evaluate=float_model_path, name_ds="validation_set")
    print('[INFO] : Evaluation complete.')
    display_figures(cfg)
 
    # Step 2: Quantize the model to INT8
    if quantization_ds:
        print('[INFO] : Using the quantization dataset to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=quantization_ds,
                                        float_model_path=float_model_path)
    else:
        print('[INFO] : Quantization dataset is not provided! Using the training set to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=train_ds,
                                        float_model_path=float_model_path)
    print('[INFO] : Quantization complete.')
 
    # Step 3: Evaluate the quantized model
    if test_ds:
        evaluate(cfg=cfg, eval_ds=test_ds, model_path_to_evaluate=quantized_model_path, name_ds="test_set")
    else:
        evaluate(cfg=cfg, eval_ds=valid_ds, model_path_to_evaluate=quantized_model_path, name_ds="validation_set")
    print('[INFO] : Evaluation complete.')
    display_figures(cfg)
 
    # Step 4: Benchmark on the target board
    benchmark(cfg=cfg, model_path_to_benchmark=quantized_model_path, credentials=credentials)
    print('[INFO] : Benchmarking complete.')
 
 
def chain_qb(cfg: DictConfig = None, float_model_path: str = None,
             train_ds: tf.data.Dataset = None,
             quantization_ds: tf.data.Dataset = None) -> None:
    """
    @brief Executes the Quantization → Benchmarking pipeline (chain_qb).
 
    @details
    Useful when accuracy evaluation is not needed and the goal is to quickly measure
    the on-device performance of a quantized model (latency, memory usage).
 
    @param cfg               Hydra configuration dictionary.
    @param float_model_path  Path to the float32 model to quantize.
    @param train_ds          Training dataset (fallback for quantization calibration).
    @param quantization_ds   Dedicated calibration dataset for INT8 quantization.
 
    @return None
    """
    credentials = None
    if cfg.tools.stm32ai.on_cloud:
        _, _, credentials = cloud_connect(stm32ai_version=cfg.tools.stm32ai.version)
 
    if quantization_ds:
        print('[INFO] : Using the quantization dataset to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=quantization_ds,
                                        float_model_path=float_model_path)
    elif train_ds:
        print('[INFO] : Quantization dataset is not provided! Using the training set to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=train_ds,
                                        float_model_path=float_model_path)
    else:
        print('[INFO] : Neither quantization dataset nor training set are provided! '
              'Using fake data to quantize the model. The model performance will not be accurate.')
        quantized_model_path = quantize(cfg=cfg, fake=True)
    print('[INFO] : Quantization complete.')
 
    benchmark(cfg=cfg, model_path_to_benchmark=quantized_model_path, credentials=credentials)
    print('[INFO] : Benchmarking complete.')
 
 
def chain_eqe(cfg: DictConfig = None, float_model_path: str = None,
              train_ds: tf.data.Dataset = None,
              valid_ds: tf.data.Dataset = None,
              quantization_ds: tf.data.Dataset = None,
              test_ds: tf.data.Dataset = None) -> None:
    """
    @brief Executes the Evaluation → Quantization → Evaluation pipeline (chain_eqe).
 
    @details
    Evaluates accuracy before and after INT8 quantization to measure the accuracy
    degradation introduced by the quantization process. No on-device benchmarking.
 
    @param cfg               Hydra configuration dictionary.
    @param float_model_path  Path to the float32 model.
    @param train_ds          Training dataset (fallback for quantization calibration).
    @param valid_ds          Validation dataset for evaluation.
    @param quantization_ds   Dedicated calibration dataset for INT8 quantization.
    @param test_ds           Test dataset (takes priority over valid_ds).
 
    @return None
    """
    # Evaluate float model
    if test_ds:
        evaluate(cfg=cfg, eval_ds=test_ds, model_path_to_evaluate=float_model_path, name_ds="test_set")
    else:
        evaluate(cfg=cfg, eval_ds=valid_ds, model_path_to_evaluate=float_model_path, name_ds="validation_set")
    print('[INFO] : Evaluation complete.')
    display_figures(cfg)
 
    # Quantize
    if quantization_ds:
        print('[INFO] : Using the quantization dataset to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=quantization_ds,
                                        float_model_path=float_model_path)
    else:
        print('[INFO] : Quantization dataset is not provided! Using the training set to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=train_ds,
                                        float_model_path=float_model_path)
    print('[INFO] : Quantization complete.')
 
    # Evaluate quantized model
    if test_ds:
        evaluate(cfg=cfg, eval_ds=test_ds, model_path_to_evaluate=quantized_model_path, name_ds="test_set")
    else:
        evaluate(cfg=cfg, eval_ds=valid_ds, model_path_to_evaluate=quantized_model_path, name_ds="validation_set")
    print('[INFO] : Evaluation complete.')
    display_figures(cfg)
 
 
def chain_tqeb(cfg: DictConfig = None, train_ds: tf.data.Dataset = None,
               valid_ds: tf.data.Dataset = None,
               quantization_ds: tf.data.Dataset = None,
               test_ds: tf.data.Dataset = None) -> None:
    """
    @brief Executes the full Training → Quantization → Evaluation → Benchmarking pipeline (chain_tqeb).
 
    @details
    This is the most complete pipeline, covering the entire model lifecycle from training
    to on-device performance measurement. It is particularly useful when starting from scratch
    or when fine-tuning a model for a new dataset.
 
    Pipeline steps:
      1. Train the model on the provided training dataset.
      2. Quantize the trained model to INT8.
      3. Evaluate the quantized model for accuracy.
      4. Benchmark on the target STM32 board.
 
    @param cfg             Hydra configuration dictionary.
    @param train_ds        Training dataset.
    @param valid_ds        Validation dataset.
    @param quantization_ds Dedicated calibration dataset (falls back to train_ds if not provided).
    @param test_ds         Test dataset (takes priority over valid_ds for evaluation).
 
    @return None
    """
    credentials = None
    if cfg.tools.stm32ai.on_cloud:
        _, _, credentials = cloud_connect(stm32ai_version=cfg.tools.stm32ai.version)
 
    # Step 1: Train
    if test_ds:
        trained_model_path = train(cfg=cfg, train_ds=train_ds, valid_ds=valid_ds, test_ds=test_ds)
    else:
        trained_model_path = train(cfg=cfg, train_ds=train_ds, valid_ds=valid_ds)
    print('[INFO] : Training complete.')
 
    # Step 2: Quantize
    if quantization_ds:
        print('[INFO] : Using the quantization dataset to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=quantization_ds,
                                        float_model_path=trained_model_path)
    else:
        print('[INFO] : Quantization dataset is not provided! Using the training set to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=train_ds,
                                        float_model_path=trained_model_path)
    print('[INFO] : Quantization complete.')
 
    # Step 3: Evaluate quantized model
    if test_ds:
        evaluate(cfg=cfg, eval_ds=test_ds, model_path_to_evaluate=quantized_model_path, name_ds="test_set")
    else:
        evaluate(cfg=cfg, eval_ds=valid_ds, model_path_to_evaluate=quantized_model_path, name_ds="validation_set")
    print('[INFO] : Evaluation complete.')
    display_figures(cfg)
 
    # Step 4: Benchmark on target board
    benchmark(cfg=cfg, model_path_to_benchmark=quantized_model_path, credentials=credentials)
    print('[INFO] : Benchmarking complete.')
 
 
def chain_tqe(cfg: DictConfig = None, train_ds: tf.data.Dataset = None,
              valid_ds: tf.data.Dataset = None,
              quantization_ds: tf.data.Dataset = None,
              test_ds: tf.data.Dataset = None) -> None:
    """
    @brief Executes the Training → Quantization → Evaluation pipeline (chain_tqe).
 
    @details
    Similar to chain_tqeb but without the final on-device benchmarking step.
    Useful when the goal is to verify accuracy after quantization without needing
    to connect a physical STM32 board.
 
    @param cfg             Hydra configuration dictionary.
    @param train_ds        Training dataset.
    @param valid_ds        Validation dataset.
    @param quantization_ds Dedicated calibration dataset (falls back to train_ds if not provided).
    @param test_ds         Test dataset (takes priority over valid_ds for evaluation).
 
    @return None
    """
    if test_ds:
        trained_model_path = train(cfg=cfg, train_ds=train_ds, valid_ds=valid_ds, test_ds=test_ds)
    else:
        trained_model_path = train(cfg=cfg, train_ds=train_ds, valid_ds=valid_ds)
    print('[INFO] : Training complete.')
 
    if quantization_ds:
        print('[INFO] : Using the quantization dataset to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=quantization_ds,
                                        float_model_path=trained_model_path)
    else:
        print('[INFO] : Quantization dataset is not provided! Using the training set to quantize the model.')
        quantized_model_path = quantize(cfg=cfg, quantization_ds=train_ds,
                                        float_model_path=trained_model_path)
    print('[INFO] : Quantization complete.')
 
    if test_ds:
        evaluate(cfg=cfg, eval_ds=test_ds, model_path_to_evaluate=quantized_model_path, name_ds="test_set")
    else:
        evaluate(cfg=cfg, eval_ds=valid_ds, model_path_to_evaluate=quantized_model_path, name_ds="validation_set")
    print('[INFO] : Evaluation complete.')
    display_figures(cfg)
 
 
def process_mode(mode: str = None,
                 configs: DictConfig = None,
                 train_ds: tf.data.Dataset = None,
                 valid_ds: tf.data.Dataset = None,
                 quantization_ds: tf.data.Dataset = None,
                 test_ds: tf.data.Dataset = None,
                 float_model_path: Optional[str] = None,
                 fake: Optional[bool] = False) -> None:
    """
    @brief Dispatches execution to the appropriate pipeline based on the operation mode.
 
    @details
    This function acts as a central dispatcher. It reads the `operation_mode` field
    from the configuration and calls the corresponding function or chain.
 
    Supported modes:
      - 'training'     : Train a model.
      - 'evaluation'   : Evaluate model accuracy on a dataset.
      - 'quantization' : Quantize a float model to INT8.
      - 'deployment'   : Deploy the model onto the STM32 board (generates C code, compiles, flashes).
      - 'prediction'   : Run inference on new input data.
      - 'benchmarking' : Measure on-device performance metrics.
      - 'chain_tqeb'   : Training → Quantization → Evaluation → Benchmarking.
      - 'chain_tqe'    : Training → Quantization → Evaluation.
      - 'chain_eqe'    : Evaluation → Quantization → Evaluation.
      - 'chain_qb'     : Quantization → Benchmarking.
      - 'chain_eqeb'   : Evaluation → Quantization → Evaluation → Benchmarking.
      - 'chain_qd'     : Quantization → Deployment.
 
    @note In deployment mode for STM32N6570-DK, after flashing the user must manually
          toggle the boot switches and power-cycle the board.
 
    @param mode            Operation mode string (e.g., 'deployment', 'chain_qd').
    @param configs         Hydra configuration dictionary.
    @param train_ds        Training TensorFlow dataset.
    @param valid_ds        Validation TensorFlow dataset.
    @param quantization_ds Calibration dataset for INT8 quantization.
    @param test_ds         Test TensorFlow dataset.
    @param float_model_path Path to the float32 model file.
    @param fake            If True, use synthetic data for quantization calibration.
 
    @return None
    @throws ValueError if an unsupported operation_mode is provided.
    """
    mlflow.log_param("model_path", configs.general.model_path)
    log_to_file(configs.output_dir, f'operation_mode: {mode}')
 
    if mode == 'training':
        if test_ds:
            train(cfg=configs, train_ds=train_ds, valid_ds=valid_ds, test_ds=test_ds)
        else:
            train(cfg=configs, train_ds=train_ds, valid_ds=valid_ds)
        display_figures(configs)
        print('[INFO] : Training complete.')
 
    elif mode == 'evaluation':
        # Generate and load the model on the STM32N6 via stedgeai, then validate on device
        gen_load_val(cfg=configs)
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        if test_ds:
            evaluate(cfg=configs, eval_ds=test_ds, name_ds="test_set")
        else:
            evaluate(cfg=configs, eval_ds=valid_ds, name_ds="validation_set")
        display_figures(configs)
        print('[INFO] : Evaluation complete.')
 
    elif mode == 'deployment':
        # Select MPU or MCU deployment path based on hardware_type config field
        if configs.hardware_type == "MPU":
            deploy_mpu(cfg=configs)
        else:
            deploy(cfg=configs)
        print('[INFO] : Deployment complete.')
        if configs.deployment.hardware_setup.board == "STM32N6570-DK":
            print('[INFO] : On STM32N6570-DK, please toggle the boot switches to the left and power cycle the board.')
 
    elif mode == 'quantization':
        # Select calibration data source
        if quantization_ds:
            input_ds = quantization_ds
            fake = False
            print('[INFO] : Using the quantization dataset to quantize the model.')
        elif train_ds:
            print('[INFO] : Quantization dataset is not provided! Using the training set to quantize the model.')
            input_ds = train_ds
            fake = False
        else:
            input_ds = None
            fake = True
            print('[INFO] : Neither quantization dataset nor training set are provided! '
                  'Using fake data to quantize the model. The model performance will not be accurate.')
        quantize(cfg=configs, quantization_ds=input_ds, fake=fake)
        print('[INFO] : Quantization complete.')
 
    elif mode == 'prediction':
        # Load model on device, then run prediction pipeline
        gen_load_val_predict(cfg=configs)
        os.chdir(os.path.dirname(os.path.realpath(__file__)))
        predict(cfg=configs)
        print('[INFO] : Prediction complete.')
 
    elif mode == 'benchmarking':
        benchmark(cfg=configs)
        print('[INFO] : Benchmark complete.')
 
    elif mode == 'chain_tqeb':
        chain_tqeb(cfg=configs, train_ds=train_ds, valid_ds=valid_ds,
                   quantization_ds=quantization_ds, test_ds=test_ds)
        print('[INFO] : chain_tqeb complete.')
 
    elif mode == 'chain_tqe':
        chain_tqe(cfg=configs, train_ds=train_ds, valid_ds=valid_ds,
                  quantization_ds=quantization_ds, test_ds=test_ds)
        print('[INFO] : chain_tqe complete.')
 
    elif mode == 'chain_eqe':
        chain_eqe(cfg=configs, float_model_path=float_model_path, train_ds=train_ds,
                  valid_ds=valid_ds, quantization_ds=quantization_ds, test_ds=test_ds)
        print('[INFO] : chain_eqe complete.')
 
    elif mode == 'chain_qb':
        chain_qb(cfg=configs, float_model_path=float_model_path, train_ds=train_ds,
                 quantization_ds=quantization_ds)
        print('[INFO] : chain_qb complete.')
 
    elif mode == 'chain_eqeb':
        chain_eqeb(cfg=configs, float_model_path=float_model_path, train_ds=train_ds,
                   valid_ds=valid_ds, quantization_ds=quantization_ds, test_ds=test_ds)
        print('[INFO] : chain_eqeb complete.')
 
    elif mode == 'chain_qd':
        chain_qd(cfg=configs, float_model_path=float_model_path, train_ds=train_ds,
                 quantization_ds=quantization_ds)
        print('[INFO] : chain_qd complete.')
 
    else:
        raise ValueError(f"Invalid mode: {mode}")
 
    # Log all output artifacts and configuration to MLflow
    mlflow.log_artifact(configs.output_dir)
    if mode in ['benchmarking', 'chain_qb', 'chain_eqeb', 'chain_tqeb']:
        mlflow.log_param("stm32ai_version", configs.tools.stm32ai.version)
        mlflow.log_param("target", configs.benchmarking.board)
    log_to_file(configs.output_dir, f'operation finished: {mode}')
 
    # Optional ClearML task logging
    if get_active_config_file() is not None:
        print(f"[INFO] : ClearML task connection")
        task = Task.current_task()
        task.connect(configs)
 
 
@hydra.main(version_base=None, config_path="", config_name="user_config")
def main(cfg: DictConfig) -> None:
    """
    @brief Main entry point of the STM32AI Model Zoo Services script.
 
    @details
    This function is decorated with @hydra.main, which means Hydra automatically
    loads the configuration from `user_config.yaml` and passes it as a DictConfig object.
 
    Execution flow:
      1. Configure GPU memory limits (if specified in the config).
      2. Parse and validate the full configuration via get_config().
      3. Initialize MLflow experiment tracking.
      4. Optionally initialize ClearML task tracking.
      5. Set the global random seed for reproducibility.
      6. Load and preprocess datasets (if required by the selected mode).
      7. Dispatch to process_mode() based on cfg.operation_mode.
 
    @note The operation mode is read from the YAML field `operation_mode`.
          Modes requiring datasets (training, evaluation, etc.) will call preprocess()
          to load and prepare the data. Modes like deployment do not require datasets.
 
    @param cfg  Hydra DictConfig object automatically populated from user_config.yaml.
 
    @return None
    """
    # Configure GPU memory ceiling to avoid OOM errors during training/quantization
    if "general" in cfg and cfg.general:
        if "gpu_memory_limit" in cfg.general and cfg.general.gpu_memory_limit:
            set_gpu_memory_limit(cfg.general.gpu_memory_limit)
            print(f"[INFO] Setting upper limit of usable GPU memory to {int(cfg.general.gpu_memory_limit)}GBytes.")
        else:
            print("[WARNING] The usable GPU memory is unlimited.\n"
                  "Please consider setting the 'gpu_memory_limit' attribute "
                  "in the 'general' section of your configuration file.")
 
    # Parse and validate the configuration file
    cfg = get_config(cfg)
    cfg.output_dir = HydraConfig.get().run.dir
    mlflow_ini(cfg)
 
    # ClearML initialization (optional — only if a valid config file exists)
    print(f"[INFO] : ClearML config check")
    if get_active_config_file() is not None:
        print(f"[INFO] : ClearML initialization and configuration")
        task = Task.init(project_name=cfg.general.project_name, task_name='pe_modelzoo_task')
        task.connect_configuration(name=cfg.operation_mode, configuration=cfg)
 
    # Set global random seed for reproducibility across TensorFlow operations
    seed = get_random_seed(cfg)
    print(f'[INFO] : The random seed for this simulation is {seed}')
    if seed is not None:
        tf.keras.utils.set_random_seed(seed)
 
    # Dispatch based on operation mode
    mode = cfg.operation_mode
 
    
    valid_modes = ['training', 'evaluation', 'chain_tqeb', 'chain_tqe']
    if mode in valid_modes:
        preprocess_output = preprocess(cfg=cfg)
        train_ds, valid_ds, quantization_ds, test_ds = preprocess_output
        process_mode(mode=mode, configs=cfg, train_ds=train_ds, valid_ds=valid_ds,
                     quantization_ds=quantization_ds, test_ds=test_ds)
 
    elif mode == 'quantization':
        
        if cfg.dataset.training_path or cfg.dataset.quantization_path:
            preprocess_output = preprocess(cfg=cfg)
            train_ds, valid_ds, quantization_ds, test_ds = preprocess_output
            process_mode(mode=mode, configs=cfg, train_ds=train_ds, valid_ds=valid_ds,
                         quantization_ds=quantization_ds, test_ds=test_ds)
        else:
            process_mode(mode=mode, configs=cfg, fake=True)
 
    else:
        
        if mode in ['chain_eqe', 'chain_qb', 'chain_eqeb', 'chain_qd']:
            if cfg.dataset.training_path or cfg.dataset.quantization_path:
                preprocess_output = preprocess(cfg=cfg)
                train_ds, valid_ds, quantization_ds, test_ds = preprocess_output
            else:
                train_ds = valid_ds = quantization_ds = test_ds = None
            process_mode(mode=mode, configs=cfg, train_ds=train_ds, valid_ds=valid_ds,
                         quantization_ds=quantization_ds, test_ds=test_ds,
                         float_model_path=cfg.general.model_path)
        else:
            
            process_mode(mode=mode, configs=cfg)
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--config-path', type=str, default='',
                        help='Path to folder containing configuration file')
    parser.add_argument('--config-name', type=str, default='user_config',
                        help='Name of the configuration file')
    parser.add_argument('params', nargs='*',
                        help='List of parameters to override in user_config.yaml')
    args = parser.parse_args()
 
    main()
 
    mlflow.log_param('config_path', args.config_path)
    mlflow.log_param('config_name', args.config_name)
    mlflow.end_run()