[FastBiasCorrection] Support models without the batch dimension

[daniil-lyakhov]

🚀 Feature request

Currently FastBiasCorrection algorithm does not support models without the batch dimension:
Reproducer:

From the branch:
https://github.com/daniil-lyakhov/nncf/tree/dl/fbc_no_batch_repro

python -m pytest /home/dlyakhov/Projects/nncf/tests/openvino/native/test_fast_bias_correction.py -k OneDimMM

Output:

    @pytest.mark.parametrize(
        "model_cls, ref_bias",
        (
            (ConvTestModel, [0.0288348, 1.0838453]),
            (ConvBNTestModel, [0.08396978, 1.1676897]),
            (OneDimMM, [0.95773065, 1.3218939, 0.81694865]),
        ),
    )
    def test_update_bias(self, model_cls, ref_bias, tmpdir):
        model = self.backend_specific_model(model_cls(), tmpdir)
        dataset = get_static_dataset(model_cls.INPUT_SIZE, self.get_transform_fn(), self.fn_to_type)
    
        quantization_algorithm = self.get_quantization_algorithm()
        graph = NNCFGraphFactory.create(model)
>       quantized_model = quantization_algorithm.apply(model, graph, dataset=dataset)

tests/cross_fw/test_templates/test_fast_bias_correction.py:121: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
nncf/quantization/algorithms/post_training/algorithm.py:113: in apply
    return self._pipeline.run_from_step(model, dataset, graph, 0, step_index_to_statistics)
nncf/quantization/algorithms/pipeline.py:164: in run_from_step
    step_model = self.run_step(step_index, step_statistics, step_model, step_graph)
nncf/quantization/algorithms/pipeline.py:119: in run_step
    current_model = pipeline_step[-1].apply(current_model, current_graph, step_statistics)
nncf/quantization/algorithms/fast_bias_correction/algorithm.py:182: in apply
    input_blob = self._backend_entity.create_input_data(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

shape = [3], data = [nncf.Tensor(0.6222554)], input_name = 'Parameter_/linear/MatMul.0', channel_axis = 0

    @staticmethod
    def create_input_data(
        shape: tuple[int], data: list[Tensor], input_name: str, channel_axis: int
    ) -> dict[str, np.ndarray]:
        blob = np.zeros(shape, dtype=data[0].data.dtype)
        for j, idx in enumerate(np.ndindex(blob.shape[channel_axis])):
            index = tuple(slice(None) if i != channel_axis else idx for i in range(blob.ndim))
>           blob[index] = data[j].data
E           IndexError: list index out of range

The task is to support models without the batch dimension in the FastBiasCorrection

Feature Use Case

To be able to use FastBiasCorrection model with models without the batch dimension

Are you going to submit a PR?

• Yes I'd like to help by submitting a PR!

[Thakor-Yashpal]

Hi, thanks for filing this feature request with the detailed reproducer. I can take on this issue and work on submitting a PR to add support for models without a batch dimension in the FastBiasCorrection algorithm.

[daniil-lyakhov]

@Thakor-Yashpal, hello! Sure, you can try to fix it! I have to warn you that this task is not intended as a GFI, so the solution for this problem will require deep dive to the NNCF architecture and possibly a non trivial fix. If you ready to accept such challenge - I'll help you as much as I can!

[darshil929]

Hi @daniil-lyakhov, I would like to work on this one. I reproduced it on current develop by adding the OneDimMM case to the FBC template (the model class is already in tree from #3479).

Two observations. First, the failure point has moved since the issue was filed: the crash now happens earlier, in _get_fp_outputs, since the tensor iterator now raises TypeError: iteration over a 0-d tensor on the scalar mean. With the same atleast_1d treatment you applied to BC in #3479, the original IndexError in create_input_data appears next. The torch backend fails the same way in _get_fp_inputs. Second, the root cause seems to sit in statistics collection: get_mean_statistic_collector treats channel_axis == 0 as the batch axis and picks the batch mean reducer, so a 1D activation gets averaged over its only axis and the per channel statistics collapse to a scalar.

My plan is rank aware reducer selection where the collectors are built (the static shape is available in get_statistic_points), atleast_1d parity with BC, and checking the create_input_data contract per backend, validated against the OneDimMM reference biases already in tree. Does this direction match what you had in mind?