IndexError: index 1 is out of bounds for axis 1 with size 1

[ueno-hiroshi002]

Issue Type

Bug

Source

pip (model-compression-toolkit)

MCT Version

2.6.0

OS Platform and Distribution

Linux Ubuntu 22.04

Python version

3.10

Describe the issue

I tried post-training quantization on a EfficientViT(m) model from timm(1.0.24).

from timm import create_model
float_model = create_model("efficientvit_m0.r224_in1k", pretrained=True)

However, I got the following error:

IndexError: index 1 is out of bounds for axis 1 with size 1

The following model causes the same error:

create_model("shvit_s1.in1k", pretrained=True)

Expected behaviour

Success PTQ.

Code to reproduce the issue

import numpy as np
from timm import create_model
from timm.data.transforms_factory import create_transform
from timm.data import resolve_data_config
import model_compression_toolkit as mct

# Prepare Model
float_model = create_model("efficientvit_m0.r224_in1k", pretrained=True)
transform = create_transform(**resolve_data_config(float_model.pretrained_cfg, model=float_model), is_training=False)

# Prepare Data
def representative_data_gen():
    yield [np.random.random((1, 3, 224, 224))]

# MCT Quantization
quantized_model, quantization_info = mct.ptq.pytorch_post_training_quantization(
                in_module=float_model,
                representative_data_gen=representative_data_gen)

Log output

Traceback (most recent call last):
  File "/home/psnrdu/sss/model_augment/minimum_efficientvit_m.py", line 16, in <module>
    quantized_model, quantization_info = mct.ptq.pytorch_post_training_quantization(
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/ptq/pytorch/quantization_facade.py", line 123, in pytorch_post_training_quantization
    tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_module,
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/runner.py", line 99, in core_runner
    graph = graph_preparation_runner(in_model,
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/graph_prep_runner.py", line 82, in graph_preparation_runner
    transformed_graph = get_finalized_graph(graph,
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/graph_prep_runner.py", line 148, in get_finalized_graph
    transformed_graph = substitute(transformed_graph, fw_impl.get_residual_collapsing_substitution())
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/common/substitutions/apply_substitutions.py", line 36, in substitute
    graph = substitution.substitute(graph, idn)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/common/substitutions/residual_collapsing.py", line 82, in substitute
    kernel_collapsed = self.residual_collapsing_fn(first_node, self.kernel_str)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/residual_collapsing.py", line 58, in residual_collapsing_fn
    kernel[i, i, idxH, idxW] += 1
IndexError: index 1 is out of bounds for axis 1 with size 1