Batched wave2vec fix

nncf/experimental/common/pruning/operations.py

@@ -408,7 +408,7 @@ def mask_propagation(
         mask = masks[0]
         output_mask = PropagationMask()
         if mask is not None and node.layer_attributes:
-            in_map, _, mode = cls.parse_reshape(node.layer_attributes)
+            in_map, out_map, mode = cls.parse_reshape(node.layer_attributes)
             input_shape = node.layer_attributes.input_shape
             output_shape = node.layer_attributes.output_shape
 
@@ -440,14 +440,26 @@ def mask_propagation(
                         group = groups[0]
                         if group.has_children():
                             raise NotImplementedError("Splitting BlockGroup with children is not implemented yet")
-                        child_groups = cls._split_group(group, input_channels, list_output_channels)
-                        for child_group, in_dim in zip(child_groups, in_map[dim]):
-                            output_mask.dim_groups_map[in_dim] = [child_group]
+                        if group.is_mixed:
+                            nncf_logger.warning("Found shrink-extend reshape pattern")
+                            # TODO: hardcode
+                            shifted_dim = in_map[dim][1]
+                            output_mask.dim_groups_map[shifted_dim] = groups
+                        else:
+                            child_groups = cls._split_group(group, input_channels, list_output_channels)
+                            for child_group, in_dim in zip(child_groups, in_map[dim]):
+                                output_mask.dim_groups_map[in_dim] = [child_group]
             elif mode == ReshapeMode.SHRINK:
                 # shrinking like: [A,B,C,D] -> S, just combine all groups under the same dimension
                 grouping = defaultdict(list)
+                in_pruning_dims = list(mask.dim_groups_map.keys())
                 for in_idx, groups in mask.dim_groups_map.items():
                     assert len(in_map[in_idx]) == 1, "assume a mapping to single int by definition of shrink"
+                    out_idx = in_map[in_idx][0]
+                    is_mixed = any(idx not in in_pruning_dims for idx in out_map[out_idx])
+                    if is_mixed:
+                        for group in groups:
+                            group.is_mixed = True
                     grouping[in_map[in_idx][0]].extend(groups)
                 output_mask.dim_groups_map = dict(grouping)
 
@@ -660,7 +672,9 @@ def _split_group(
         new_groups: List[PropagationGroup] = []
         new_blocks = cls._split_block(group.block, list_output_channels)
         for block in new_blocks:
-            new_group = PropagationGroup(block=block, producers=group.get_producers(), consumers=group.get_consumers())
+            new_group = PropagationGroup(
+                block=block, producers=group.get_producers(), consumers=group.get_consumers(), is_mixed=group.is_mixed
+            )
             group.add_child(new_group)
             new_groups.append(new_group)
         return new_groups

nncf/experimental/common/pruning/propagation_data.py

@@ -147,12 +147,14 @@ def __init__(
         block: PruningBlock,
         producers: Optional[Set[ProducerInfo]] = None,
         consumers: Optional[Set[ConsumerInfo]] = None,
+        is_mixed: bool = False,
     ) -> None:
         self.block = block
         self._children: List["PropagationGroup"] = []
         self._is_invalid = False
         self._producers = set() if producers is None else producers
         self._consumers = set() if consumers is None else consumers
+        self.is_mixed = is_mixed
 
     @property
     def is_invalid(self):
@@ -190,10 +192,12 @@ def join_groups(*args: "PropagationGroup") -> "PropagationGroup":
         )
 
         new_group = PropagationGroup(block=first_block)
+        is_mixed = any(group.is_mixed for group in args)
         for group in args:
             group.add_child(new_group)
             new_group.add_producers(group.get_producers())
             new_group.add_consumers(group.get_consumers())
+        new_group.is_mixed = is_mixed
         return new_group
 
     def add_consumers(self, consumers: Set[ConsumerInfo]):

tests/torch/pruning/experimental/test_nodes_grouping.py

@@ -426,7 +426,7 @@ def __str__(self) -> str:
     GroupTestDesc(
         model_desc=GeneralModelDesc(
             model_name="Wave2Vec 2.0",
-            input_info=dict(sample_size=[1, 400]),
+            input_info=dict(sample_size=[3, 400]),
             model_builder=partial(
                 AutoModelForAudioClassification.from_config,
                 Wav2Vec2Config(

tests/torch/sparsity/movement/helpers/run_recipe.py

@@ -280,6 +280,27 @@ def _create_model(self) -> torch.nn.Module:
         return AutoModelForAudioClassification.from_config(self.model_config)
 
 
+class Wav2Vec2RunRecipeBatched(Wav2Vec2RunRecipe):
+    default_model_config = Wav2Vec2Config(
+        vocab_size=2,
+        hidden_size=4,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        conv_dim=(4, 4),
+        conv_stride=(1, 1),
+        conv_kernel=(3, 3),
+        num_conv_pos_embeddings=3,
+        num_conv_pos_embedding_groups=1,
+        proj_codevector_dim=4,
+        classifier_proj_size=3,
+        num_labels=2,
+    )
+
+    @property
+    def model_input_info(self) -> List[ModelInputInfo]:
+        return [ModelInputInfo(shape=[3, 32], keyword="input_values")]
+
+
 class BertRunRecipe(BaseMockRunRecipe):
     model_family = "huggingface_bert"
     supports_structured_masking = True

tests/torch/sparsity/movement/test_model_saving.py

@@ -44,6 +44,7 @@
 from tests.torch.sparsity.movement.helpers import build_compression_trainer
 from tests.torch.sparsity.movement.helpers import force_update_sparsifier_binary_masks_by_threshold
 from tests.torch.sparsity.movement.helpers import initialize_sparsifier_parameters_by_linspace
+from tests.torch.sparsity.movement.helpers.run_recipe import Wav2Vec2RunRecipeBatched
 
 
 class TestONNXExport:
@@ -176,6 +177,11 @@ def test_same_outputs_in_torch_and_exported_onnx(self, tmp_path: Path, recipe: B
                     classifier_proj_size=3,
                 ),
             ),
+            Dict(
+                nncf_weight_ratio=0.74,
+                ov_weight_ratio=0.74,
+                recipe=Wav2Vec2RunRecipeBatched().model_config_(),
+            ),
             Dict(
                 nncf_weight_ratio=0.43,
                 ov_weight_ratio=0.29,
@@ -239,7 +245,7 @@ def test_ngraph_pruning(self, tmp_path: Path, desc: dict):
         ctrl.resolve_structured_mask()
         ctrl.populate_structured_mask()
         compression_rate = ctrl.statistics().movement_sparsity.model_statistics.sparsity_level
-
+        print(ctrl.statistics().to_str())
         onnx_model_path = str(tmp_path / "structured.onnx")
         ctrl.export_model(onnx_model_path)
         core = Core()