[PT] BatchNorm adaptation (#3726)

### Changes

Add `nncf.batch_norm_adaptation` function.
Add mode `bn_adaptaion` in example


### Related tickets

174483

### Tests


https://github.com/openvinotoolkit/nncf/actions/runs/19162855749/job/54776694802

---------

Co-authored-by: Lyalyushkin Nikolay <[email protected]>

Author: AlexanderDokuchaev

examples/pruning/torch/resnet18/README.md

@@ -31,7 +31,12 @@ python3 -m pip install ../../../../ -r requirements.txt
 It's pretty simple. The example does not require additional preparation. It will do the preparation itself, such as loading the dataset and model, etc.
 
 ```bash
+# To run Magnitude-Based pruning
 python main.py
-# Or to run Regularization-Based pruning
+
+# To run Magnitude-Based pruning with batch norm adaptation
+python main.py --mode mag_bn
+
+# To run Regularization-Based pruning
 python main.py --mode rb
 ```

examples/pruning/torch/resnet18/main.py

@@ -12,6 +12,7 @@
 import os
 import warnings
 from argparse import ArgumentParser
+from argparse import RawTextHelpFormatter
 from pathlib import Path
 
 import openvino as ov
@@ -48,13 +49,18 @@
 
 
 def get_argument_parser() -> ArgumentParser:
-    parser = ArgumentParser()
+    parser = ArgumentParser(formatter_class=RawTextHelpFormatter)
     parser.add_argument(
         "--mode",
         type=str,
-        choices=["magnitude", "rb"],
-        default="magnitude",
-        help="Pruning mode to use. Choices are: magnitude, rb. Default is magnitude.",
+        choices=["mag", "mag_bn", "rb"],
+        default="mag",
+        help=(
+            "Pruning mode to use. Choices are:\n"
+            " - mag: Magnitude-based pruning with fine-tuning (default).\n"
+            " - mag_bn: Magnitude-based pruning with BatchNorm adaptation without fine-tuning.\n"
+            " - rb: Regularization-based pruning with fine-tuning.\n"
+        ),
     )
     return parser
 
@@ -82,13 +88,12 @@ def get_resnet18_model(device: torch.device) -> nn.Module:
 def train_epoch(
     train_loader: DataLoader,
     model: nn.Module,
-    criterion: nn.Module,
     rb_loss: RBLoss,
     optimizer: torch.optim.Optimizer,
     device: torch.device,
 ):
-    # Switch to train mode.
     model.train()
+    criterion = nn.CrossEntropyLoss().to(device)
 
     for images, target in track(train_loader, total=len(train_loader), description="Fine tuning:"):
         images = images.to(device)
@@ -107,7 +112,6 @@ def train_epoch(
 
 @torch.no_grad()
 def validate(val_loader: torch.utils.data.DataLoader, model: torch.nn.Module, device: torch.device) -> float:
-    # Switch to evaluate mode.
     model.eval()
 
     correct = 0
@@ -201,14 +205,20 @@ def main() -> float:
 
     ###############################################################################
     # Step 2: Prune model
-    print(os.linesep + "[Step 2]: Prune model and specify training parameters")
+    print(os.linesep + "[Step 2] Prune model and specify training parameters")
 
-    if pruning_mode == "magnitude":
+    if pruning_mode == "mag_bn":
+        pruned_model = nncf.prune(
+            model,
+            mode=PruneMode.UNSTRUCTURED_MAGNITUDE_GLOBAL,
+            ratio=0.6,
+            examples_inputs=example_input,
+        )
+    elif pruning_mode == "mag":
         pruned_model = nncf.prune(
             model,
             mode=PruneMode.UNSTRUCTURED_MAGNITUDE_GLOBAL,
             ratio=0.7,
-            ignored_scope=nncf.IgnoredScope(),
             examples_inputs=example_input,
         )
         num_epochs = 2
@@ -217,11 +227,10 @@ def main() -> float:
             model=model, mode=PruneMode.UNSTRUCTURED_MAGNITUDE_GLOBAL, steps={0: 0.5, 1: 0.7}
         )
         optimizer = torch.optim.Adam(pruned_model.parameters(), lr=1e-5)
-    else:
+    elif pruning_mode == "rb":
         pruned_model = nncf.prune(
             model,
             mode=PruneMode.UNSTRUCTURED_REGULARIZATION_BASED,
-            ignored_scope=nncf.IgnoredScope(),
             examples_inputs=example_input,
         )
         num_epochs = 30
@@ -237,32 +246,52 @@ def main() -> float:
                 {"params": mask_params, "lr": 1e-2, "weight_decay": 0.0},
             ]
         )
-
-    criterion = nn.CrossEntropyLoss().to(device)
+    else:
+        msg = f"Unsupported pruning mode: {pruning_mode}, please choose from ['mag', 'mag_bn', 'rb']"
+        raise ValueError(msg)
 
     ###############################################################################
     # Step 3: Fine tune
     print(os.linesep + "[Step 3] Fine tune with multi step pruning ratio scheduler")
 
-    for epoch in range(num_epochs):
-        print(os.linesep + f"Train epoch: {epoch}")
-        scheduler.step()
-        train_epoch(train_loader, pruned_model, criterion, rb_loss, optimizer, device=device)
+    if pruning_mode == "mag_bn":
+        acc1_before = validate(val_loader, pruned_model, device)
+        print(f"Accuracy@1 of pruned model before BatchNorm adaptation: {acc1_before:.3f}")
+
+        def transform_fn(batch: tuple[torch.Tensor, int]) -> torch.Tensor:
+            inputs, _ = batch
+            return inputs.to(device=device)
+
+        calibration_dataset = nncf.Dataset(train_loader, transform_func=transform_fn)
+
+        pruned_model = nncf.batch_norm_adaptation(
+            pruned_model,
+            calibration_dataset=calibration_dataset,
+            num_iterations=200,
+        )
 
         acc1 = validate(val_loader, pruned_model, device)
-        print(f"Current pruning ratio: {scheduler.current_ratio:.3f}")
-        print(f"Accuracy@1 of pruned model after {epoch} epoch: {acc1:.3f}")
+        print(f"Accuracy@1 of pruned model after BatchNorm adaptation: {acc1:.3f}")
+    else:
+        for epoch in range(num_epochs):
+            print(os.linesep + f"Train epoch: {epoch}")
+            scheduler.step()
+            train_epoch(train_loader, pruned_model, rb_loss, optimizer, device=device)
+
+            acc1 = validate(val_loader, pruned_model, device)
+            print(f"Current pruning ratio: {scheduler.current_ratio:.3f}")
+            print(f"Accuracy@1 of pruned model after {epoch} epoch: {acc1:.3f}")
 
     ###############################################################################
     # Step 4: Print per tensor pruning statistics
-    print(os.linesep + "[Step 4]: Pruning statistics")
+    print(os.linesep + "[Step 4] Pruning statistics")
 
     pruning_stat = nncf.pruning_statistic(pruned_model)
     print(pruning_stat)
 
     ###############################################################################
     # Step 5: Export models
-    print(os.linesep + "[Step 5]: Export models")
+    print(os.linesep + "[Step 5] Export models")
     ir_path = ROOT / f"{BASE_MODEL_NAME}_pruned.xml"
     ov_model = ov.convert_model(pruned_model.cpu(), example_input=example_input.cpu(), input=tuple(example_input.shape))
     ov.save_model(ov_model, ir_path, compress_to_fp16=False)

src/nncf/__init__.py

@@ -44,6 +44,7 @@
 from nncf.parameters import SensitivityMetric as SensitivityMetric
 from nncf.parameters import StripFormat as StripFormat
 from nncf.parameters import TargetDevice as TargetDevice
+from nncf.pruning.prune_model import batch_norm_adaptation as batch_norm_adaptation
 from nncf.pruning.prune_model import prune as prune
 from nncf.pruning.prune_model import pruning_statistic as pruning_statistic
 from nncf.quantization import QuantizationPreset as QuantizationPreset

src/nncf/pruning/prune_model.py

@@ -17,6 +17,7 @@
 from nncf.common.utils.backend import BackendType
 from nncf.common.utils.backend import get_backend
 from nncf.common.utils.helpers import create_table
+from nncf.data.dataset import Dataset
 from nncf.parameters import PruneMode
 from nncf.scopes import IgnoredScope
 
@@ -51,6 +52,29 @@ def prune(
     return model
 
 
+def batch_norm_adaptation(
+    model: TModel, calibration_dataset: Dataset, *, num_iterations: Optional[int] = None
+) -> TModel:
+    """
+    Adapt the batch normalization layers of the given model using the provided dataset.
+    This function runs a specified number of iterations through the model
+    to update the running statistics of the batch normalization layers.
+
+    :param model: The model to adapt.
+    :param calibration_dataset: The dataset to use for the adaptation.
+    :param num_iterations: The number of iterations to use for adaptation.
+        If set to None, the adaptation will run for the entire dataset.
+    """
+    backend = get_backend(model)
+    if backend == BackendType.TORCH:
+        from nncf.torch.function_hook.pruning.batch_norm_adaptation import batch_norm_adaptation
+
+        return batch_norm_adaptation(model, calibration_dataset=calibration_dataset, num_iterations=num_iterations)
+
+    msg = f"Batch norm adaptation is not supported for the {backend} backend."
+    raise nncf.InternalError(msg)
+
+
 @dataclass
 class TensorPruningStatistic:
     """

src/nncf/torch/function_hook/pruning/batch_norm_adaptation.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2025 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from contextlib import contextmanager
+from typing import Generator, Optional, TypeVar
+
+import torch
+from torch import nn
+
+from nncf import Dataset
+from nncf.common.logging.track_progress import track
+
+TModel = TypeVar("TModel", bound=nn.Module)
+
+
+@torch.no_grad()
+def batch_norm_adaptation(
+    model: TModel, calibration_dataset: Dataset, *, num_iterations: Optional[int] = None
+) -> TModel:
+    """
+    Adapt the batch normalization layers of the given model using the provided dataset.
+
+    This function runs a specified number of iterations (batches) through the model
+    to update the running statistics of the batch normalization layers.
+
+    :param model: The model to adapt.
+    :param calibration_dataset: The dataset to use for the adaptation.
+    :param num_iterations: The number of iterations (batches) to use for adaptation.
+        If set to None, the adaptation will run for the entire dataset.
+    """
+    with set_batchnorm_train_only(model):
+        total = calibration_dataset.get_length()
+        if num_iterations is not None:
+            total = min(num_iterations, total) if total is not None else num_iterations
+
+        for idx, input_data in track(
+            enumerate(calibration_dataset.get_inference_data()),
+            total=total,
+            description="Batch norm adaptation",
+        ):
+            if num_iterations is not None and idx >= num_iterations:
+                break
+
+            if isinstance(input_data, dict):
+                model(**input_data)
+            elif isinstance(input_data, tuple):
+                model(*input_data)
+            else:
+                model(input_data)
+
+    return model
+
+
+@contextmanager
+def set_batchnorm_train_only(model: nn.Module) -> Generator[None, None, None]:
+    """
+    Context manager that sets only BatchNorm modules to train mode,
+    while keeping all other modules in eval mode.
+    Restores the original training states afterward.
+
+    :param model: The model.
+    """
+    # Store the original training states
+    original_states = {}
+    for name, module in model.named_modules():
+        original_states[name] = module.training
+
+    try:
+        # Set all modules to eval, then only BN to train
+        model.eval()
+        for module in model.modules():
+            if isinstance(module, nn.modules.batchnorm._BatchNorm):
+                module.train()
+        yield
+    finally:
+        # Restore original training states
+        for name, module in model.named_modules():
+            module.train(original_states[name])

src/nncf/torch/function_hook/pruning/magnitude/algo.py

@@ -9,6 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TypeVar
+
 import torch
 from torch import nn
 
@@ -20,13 +22,15 @@
 from nncf.torch.function_hook.wrapper import register_post_function_hook
 from nncf.torch.model_graph_manager import get_const_data_by_name
 
+TModel = TypeVar("TModel", bound=nn.Module)
+
 
 def apply_magnitude_pruning(
-    model: nn.Module,
+    model: TModel,
     parameters: list[str],
     mode: PruneMode,
     ratio: float,
-) -> nn.Module:
+) -> TModel:
     """
     Prunes the specified parameters of the given model using unstructured pruning.
 

src/nncf/torch/function_hook/pruning/prune_model.py

@@ -8,7 +8,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Optional
+from typing import Any, Optional, TypeVar
 
 from torch import nn
 
@@ -25,6 +25,8 @@
 from nncf.torch.function_hook.wrapper import wrap_model
 from nncf.torch.model_graph_manager import get_const_node
 
+TModel = TypeVar("TModel", bound=nn.Module)
+
 OPERATORS_WITH_WEIGHTS_METATYPES = [
     om.PTConv1dMetatype,
     om.PTConv2dMetatype,
@@ -43,12 +45,12 @@
 
 
 def prune(
-    model: nn.Module,
+    model: TModel,
     mode: PruneMode,
     ratio: Optional[float] = None,
     ignored_scope: Optional[IgnoredScope] = None,
     examples_inputs: Optional[Any] = None,
-) -> nn.Module:
+) -> TModel:
     if examples_inputs is None:
         msg = "`sparsity` function requires `examples_inputs` argument to be specified for Torch backend"
         raise nncf.InternalError(msg)

src/nncf/torch/function_hook/pruning/rb/algo.py

@@ -9,6 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TypeVar
+
 from torch import nn
 
 import nncf
@@ -18,11 +20,13 @@
 from nncf.torch.function_hook.wrapper import register_post_function_hook
 from nncf.torch.model_graph_manager import get_const_data_by_name
 
+TModel = TypeVar("TModel", bound=nn.Module)
+
 
 def apply_regularization_based_pruning(
-    model: nn.Module,
+    model: TModel,
     parameters: list[str],
-) -> nn.Module:
+) -> TModel:
     """
     :param model: The neural network model to be pruned.
     :param parameters: A list of parameter names to be pruned.