Start of making bias correction work with Conv1d

[V0XNIHILI]

In this PR, I already fixed all the issues that appear in the Python code when trying to run bias correction on models with nn.Conv1d layers in it. I do not know how to fix this remaining issue, so if anybody from the QUIC team could help me with that, that would be amazing.

Code to replicate

from aimet_torch import bias_correction
from aimet_torch.quantsim import QuantParams

def apply_bias_correction(model: torch.nn.Module, data_loader: DataLoader):
    """
    Applies Bias-Correction on the model.
    :param model: The model to quantize
    :param evaluator: Evaluator used during quantization
    :param dataloader: DataLoader used during quantization
    :param logdir: Log directory used for storing log files
    :return: None
    """
    # Rounding mode can be 'nearest' or 'stochastic'
    rounding_mode = 'nearest'

    # Number of samples used during quantization
    num_quant_samples = 16

    # Number of samples used for bias correction
    num_bias_correct_samples = 16

    params = QuantParams(weight_bw=8, act_bw=8, round_mode=rounding_mode, quant_scheme='tf_enhanced')

    # Perform Bias Correction
    bias_correction.correct_bias(model.to(device='cpu'), params, num_quant_samples=num_quant_samples,
                                 data_loader=data_loader, num_bias_correct_samples=num_bias_correct_samples)from aimet_torch import bias_correction
from aimet_torch.quantsim import QuantParams

def apply_bias_correction(model: torch.nn.Module, data_loader: DataLoader):
    """
    Applies Bias-Correction on the model.
    :param model: The model to quantize
    :param evaluator: Evaluator used during quantization
    :param dataloader: DataLoader used during quantization
    :param logdir: Log directory used for storing log files
    :return: None
    """
    # Rounding mode can be 'nearest' or 'stochastic'
    rounding_mode = 'nearest'

    # Number of samples used during quantization
    num_quant_samples = 16

    # Number of samples used for bias correction
    num_bias_correct_samples = 16

    params = QuantParams(weight_bw=8, act_bw=8, round_mode=rounding_mode, quant_scheme='tf_enhanced')

    # Perform Bias Correction
    bias_correction.correct_bias(model.to(device='cpu'), params, num_quant_samples=num_quant_samples,
                                 data_loader=data_loader, num_bias_correct_samples=num_bias_correct_samples)

input_shape = (1, 40, 101)

model = model.eval()

# Performs BatchNorm fold, Cross layer scaling and High bias folding
equalize_model(model, input_shape)

# Model has nn.Conv1d layers in it, train_data_loader returns data of shape (batch size, 40, 101)
apply_bias_correction(model=model, data_loader=train_data_loader)

Remaining issue

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb](https://vscode-remote+ssh-002dremote-002bquechua1-002eewi-002etudelft-002enl.vscode-resource.vscode-cdn.net/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb) Cell 28 in ()
----> [1](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0) apply_bias_correction(model=model, data_loader=train_data_loader)

[/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb](https://vscode-remote+ssh-002dremote-002bquechua1-002eewi-002etudelft-002enl.vscode-resource.vscode-cdn.net/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb) Cell 28 in apply_bias_correction(model, data_loader)
     [22](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=21) params = QuantParams(weight_bw=8, act_bw=8, round_mode=rounding_mode, quant_scheme='tf_enhanced')
     [24](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=23) # Perform Bias Correction
---> [25](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=24) bias_correction.correct_bias(model.to(device='cpu'), params, num_quant_samples=num_quant_samples,
     [26](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=25)                              data_loader=data_loader, num_bias_correct_samples=num_bias_correct_samples)

File [~/anaconda3/envs/meta-learning-arena3.8/lib/python3.8/site-packages/aimet_torch/bias_correction.py:342](https://vscode-remote+ssh-002dremote-002bquechua1-002eewi-002etudelft-002enl.vscode-resource.vscode-cdn.net/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/~/anaconda3/envs/meta-learning-arena3.8/lib/python3.8/site-packages/aimet_torch/bias_correction.py:342), in correct_bias(model, quant_params, num_quant_samples, data_loader, num_bias_correct_samples, conv_bn_dict, perform_only_empirical_bias_corr, layers_to_ignore)
    339         reference_output_batch = reference_output_batch.reshape(extended_shape)
    340         quantized_model_output_batch = quantized_model_output_batch.reshape(extended_shape)
--> 342     bias_correction.storePreActivationOutput(reference_output_batch)
    343     bias_correction.storeQuantizedPreActivationOutput(quantized_model_output_batch)
    345 call_empirical_mo_correct_bias(module, bias_correction)

ValueError: array has incorrect number of dimensions: 3; expected 4---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb](https://vscode-remote+ssh-002dremote-002bquechua1-002eewi-002etudelft-002enl.vscode-resource.vscode-cdn.net/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb) Cell 28 in ()
----> [1](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0) apply_bias_correction(model=model, data_loader=train_data_loader)

[/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb](https://vscode-remote+ssh-002dremote-002bquechua1-002eewi-002etudelft-002enl.vscode-resource.vscode-cdn.net/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb) Cell 28 in apply_bias_correction(model, data_loader)
     [22](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=21) params = QuantParams(weight_bw=8, act_bw=8, round_mode=rounding_mode, quant_scheme='tf_enhanced')
     [24](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=23) # Perform Bias Correction
---> [25](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=24) bias_correction.correct_bias(model.to(device='cpu'), params, num_quant_samples=num_quant_samples,
     [26](vscode-notebook-cell://ssh-remote%2Bquechua1.ewi.tudelft.nl/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/quant_tcn_kws_aimet.ipynb#Y125sdnNjb2RlLXJlbW90ZQ%3D%3D?line=25)                              data_loader=data_loader, num_bias_correct_samples=num_bias_correct_samples)

File [~/anaconda3/envs/meta-learning-arena3.8/lib/python3.8/site-packages/aimet_torch/bias_correction.py:342](https://vscode-remote+ssh-002dremote-002bquechua1-002eewi-002etudelft-002enl.vscode-resource.vscode-cdn.net/space/ddenblanken/Projects/meta-learning-arena/src/metalarena/experiments/~/anaconda3/envs/meta-learning-arena3.8/lib/python3.8/site-packages/aimet_torch/bias_correction.py:342), in correct_bias(model, quant_params, num_quant_samples, data_loader, num_bias_correct_samples, conv_bn_dict, perform_only_empirical_bias_corr, layers_to_ignore)
    339         reference_output_batch = reference_output_batch.reshape(extended_shape)
    340         quantized_model_output_batch = quantized_model_output_batch.reshape(extended_shape)
--> 342     bias_correction.storePreActivationOutput(reference_output_batch)
    343     bias_correction.storeQuantizedPreActivationOutput(quantized_model_output_batch)
    345 call_empirical_mo_correct_bias(module, bias_correction)

ValueError: array has incorrect number of dimensions: 3; expected 4

[quic-akhobare]

Hi @V0XNIHILI .. So specifically this code needs to be updated

[quic-akhobare]

Let me walk through the changes likely needed

• the weight tensors being sent as arguments to the correctBias call on line 208, need to be 4D tensors.
• So in the case of Conv1D, the weight tensors will be 3D, so we need something like np.expand_dims(x, 3)
• Code block at line 181 needs to be updated for ConvTranspose1D

Could you try going a bit further? If you get stuck someone else can pitch in..

[quic-akhobare]

Just restarting the job - since it failed due to a known error. So we can see the changes so far are partial but don't break existing support.

[quic-mangal]

@V0XNIHILI to add to Abhi's comment, you can refer to https://github.com/quic/aimet/blob/develop/TrainingExtensions/torch/src/python/aimet_torch/batch_norm_fold.py#:~:text=with-,_expand_shape_to_4d,-(weight_tensor)%3A
for how we Conv1D is supported

[quic-bharathr]

Hi @V0XNIHILI , could you please rebase your branch with the latest "develop" branch?
https://github.com/quic/aimet/tree/develop
That should likely fix the status checks failures. Let me know if you need any help.