AIMET Integration

[kozlov721]

Purpose

Adds option to quantize trained model using various PTQ and QAT techniques.

Specification

• Added LuxonisModel.quantize method
  • Runs selected algorithms on the model weights and exports it to ONNX
  • Support for various PQT techniques:
    • Adaround
    • Batch norm folding
    • Cross-layer equalization
    • Batch norm re-estimation
    • Sequential MSE
• Added AIMETCallback
  • Runs quantization at the end of the training
• Added a new section exporter.aimet to the Config
• Various small changes necessary to accomodate the AIMET API or to simplify the integration
  • Added more standard forward method to LuxonixLightningModule
    • The original to full_forward
  • Support for loaders outputting bare tensors
    • As opposed to the expected dictionary
  • Removed metrics, losses and visuazilers from saved node modules
  • Added print_table as a required abstract method to BaseLuxonisProgressBar
  • Implemented custom __getstate__ and __setstate__ in LuxonisLightningModule

config.yaml

# Default values
exporter:
  aimet:
    active: false

    default_output_bw: 8
    default_param_bw: 8
    default_data_type: "int"
    quant_scheme: "min_max"
    config: ~

    fold_batch_norms: false
    cross_layer_equalization: false
    batch_norm_reestimation: false
    sequential_mse: false
    adaround:
      active: false
      default_num_iterations: ~
      default_reg_param: 0.01
      default_beta_range: [20, 2]
      default_warm_start: 0.2

    epochs: 20
    optimizer:
      name: "SGD"
      params:
        lr: 0.00001
    scheduler:
      name: "StepLR"
      params:
        step_size: 5
        gamma: 0.1

Dependencies & Potential Impact

None / not applicable

Deployment Plan

None / not applicable

Testing & Validation

• Testing AIMETCallback in test_callbacks
• Testing LuxonisModel.quantize for all predefined models with the full set of PTQ techniques enabled

[kozlov721]

What if we just make this a required parameter is PTQ is active?

You mean the opt-in advanced techniques? Yeah I think this would make sense. What you want to use usually also depends on the specific model so as long as the entire AIMET is opt-in I think forcing the user to explicitly specify what to use is a good idea.

I suppose eventually we will learn what works best with our predefined models and we can then enable it by default for them with the known well working parameters.