Release v1.7.1 of NNCF to master

Author: vshampor

ReleaseNotes.md

@@ -7,6 +7,11 @@ samples distributed with the code.  The samples demonstrate the usage of compres
 public models and datasets for three different use cases: Image Classification, Object Detection,
 and Semantic Segmentation.
 
+## New in Release 1.7.1:
+Bugfixes:
+- Fixed a bug with where compressed models that were supposed to return named tuples actually returned regular tuples
+- Fixed an issue with batch norm adaptation-enabled compression runs hanging in the DDP scenario
+
 ## New in Release 1.7:
 - Adjust Padding feature to support accurate execution of U4 on VPU - when setting "target_device" to "VPU", the training-time padding values for quantized convolutions will be adjusted to better reflect VPU inference process.
 - Weighted layers that are "frozen" (i.e. have requires_grad set to False at compressed model creation time) are no longer considered for compression, to better handle transfer learning cases.

examples/classification/main.py

@@ -199,6 +199,7 @@ def autoq_eval_fn(model, eval_loader):
     if config.mode.lower() == 'train':
         train(config, compression_ctrl, model, criterion, train_criterion_fn, lr_scheduler, model_name, optimizer,
               train_loader, train_sampler, val_loader, best_acc1)
+    config.mlflow.end_run()
 
 
 def train(config, compression_ctrl, model, criterion, criterion_fn, lr_scheduler, model_name, optimizer,
@@ -267,9 +268,11 @@ def get_dataset(dataset_config, config, transform, is_train):
     if dataset_config == 'imagenet':
         prefix = 'train' if is_train else 'val'
         return datasets.ImageFolder(osp.join(config.dataset_dir, prefix), transform)
+    # For testing purposes
     if dataset_config == 'mock_32x32':
-        # For testing purposes
         return MockDataset(img_size=(32, 32), transform=transform)
+    if dataset_config == 'mock_299x299':
+        return MockDataset(img_size=(299, 299), transform=transform)
     return create_cifar(config, dataset_config, is_train, transform)
 
 
@@ -287,15 +290,16 @@ def create_cifar(config, dataset_config, is_train, transform):
 def create_datasets(config):
     dataset_config = config.dataset if config.dataset is not None else 'imagenet'
     dataset_config = dataset_config.lower()
-    assert dataset_config in ['imagenet', 'cifar100', 'cifar10', 'mock_32x32'], "Unknown dataset option"
+    assert dataset_config in ['imagenet', 'cifar100', 'cifar10', 'mock_32x32', 'mock_299x299'], \
+        "Unknown dataset option"
 
     if dataset_config == 'imagenet':
         normalize = transforms.Normalize(mean=(0.485, 0.456, 0.406),
                                          std=(0.229, 0.224, 0.225))
     elif dataset_config == 'cifar100':
         normalize = transforms.Normalize(mean=(0.5071, 0.4865, 0.4409),
                                          std=(0.2673, 0.2564, 0.2761))
-    elif dataset_config in ['cifar10', 'mock_32x32']:
+    elif dataset_config in ['cifar10', 'mock_32x32', 'mock_299x299']:
         normalize = transforms.Normalize(mean=(0.5, 0.5, 0.5),
                                          std=(0.5, 0.5, 0.5))
 

examples/common/utils.py

@@ -128,6 +128,9 @@ def safe_call(self, func: str, *args, **kwargs) -> Maybe:
         return Maybe.from_value(self._get_mlflow()).bind(
             lambda obj: Maybe.from_value(getattr(obj, func)(*args, **kwargs)))
 
+    def end_run(self):
+        self.safe_call('end_run')
+
     def _is_enabled(self):
         return self.is_suitable_mode and is_main_process()
 

nncf/dynamic_graph/context.py

@@ -475,11 +475,11 @@ def _init_thread_local(self):
         tl.operator_counters = {}
         tl.node_call_tracker = {}
 
-    def register_node_call(self, node_key: str):
-        if node_key in self._thread_local.node_call_tracker:
-            self._thread_local.node_call_tracker[node_key] += 1
+    def register_node_call(self, node: NNCFNode):
+        if node.node_id in self._thread_local.node_call_tracker:
+            self._thread_local.node_call_tracker[node.node_id] += 1
         else:
-            self._thread_local.node_call_tracker[node_key] = 1
+            self._thread_local.node_call_tracker[node.node_id] = 1
 
     def reset_node_call_counters(self):
         for k, _ in self._thread_local.node_call_tracker.items():

nncf/dynamic_graph/wrappers.py

@@ -79,7 +79,7 @@ def wrapped(*args, **kwargs):
                 node = ctx.maybe_add_node(processed_input, tensor_metas, ia_op_exec_context, module_attrs)
 
             if is_debug():
-                ctx.register_node_call(ctx.graph.get_node_key_by_id(node.node_id))
+                ctx.register_node_call(node)
 
             args = tuple(processed_input.op_args)
             kwargs = processed_input.op_kwargs

nncf/initialization.py

@@ -17,7 +17,7 @@
 from nncf.structures import QuantizationRangeInitArgs
 from nncf.utils import is_tensor
 from nncf.utils import objwalk
-from nncf.utils import training_mode_switcher
+from contextlib import contextmanager
 
 
 class InitializingDataLoader:
@@ -164,39 +164,66 @@ def __init__(self, model, init_device: str, num_bn_forget_steps):
         self.num_bn_forget_steps = num_bn_forget_steps
         self.momentum_bn_forget = 0.9
         self.original_momenta_values = {}
+        self.original_training_state = {}
 
     @staticmethod
     def _apply_to_batchnorms(func):
         def func_apply_to_bns(module):
-            if isinstance(module, torch.nn.modules.batchnorm.BatchNorm2d):
+            if isinstance(module, (torch.nn.modules.batchnorm.BatchNorm1d,
+                                   torch.nn.modules.batchnorm.BatchNorm2d,
+                                   torch.nn.modules.batchnorm.BatchNorm3d)):
                 func(module)
 
         return func_apply_to_bns
 
-    def _run_model_inference(self, data_loader, num_init_steps, device):
-        num_bn_forget_steps = self.num_bn_forget_steps
+    @contextmanager
+    def _bn_training_state_switcher(self) -> None:
+        def save_original_bn_training_state(module: torch.nn.Module):
+            self.original_training_state[module] = module.training
+
+        def set_bn_training_state(module: torch.nn.Module, state: Dict[str, bool]):
+            module.training = state
+
+        def restore_original_bn_training_state(module: torch.nn.Module):
+            module.training = self.original_training_state[module]
+
+        self.model.apply(self._apply_to_batchnorms(save_original_bn_training_state))
+        self.model.apply(self._apply_to_batchnorms(partial(set_bn_training_state, state=True)))
+        try:
+            yield
+        finally:
+            self.model.apply(self._apply_to_batchnorms(restore_original_bn_training_state))
 
+    @contextmanager
+    def _bn_momentum_switcher(self) -> None:
         def set_bn_momentum(module, momentum_value):
             module.momentum = momentum_value
 
-        def save_original_bn_momenta(module):
+        def save_original_bn_momentum(module: torch.nn.Module):
             self.original_momenta_values[module] = module.momentum
 
-        def restore_original_bn_momenta(module):
+        def restore_original_bn_momentum(module: torch.nn.Module):
             module.momentum = self.original_momenta_values[module]
 
-        with training_mode_switcher(self.model, is_training=True):
-            self.model.apply(self._apply_to_batchnorms(save_original_bn_momenta))
-            self.model.apply(self._apply_to_batchnorms(partial(set_bn_momentum,
-                                                               momentum_value=self.momentum_bn_forget)))
+        self.model.apply(self._apply_to_batchnorms(save_original_bn_momentum))
+        self.model.apply(self._apply_to_batchnorms(partial(set_bn_momentum,
+                                                           momentum_value=self.momentum_bn_forget)))
+        try:
+            yield
+        finally:
+            self.model.apply(self._apply_to_batchnorms(restore_original_bn_momentum))
 
-            for i, loaded_item in enumerate(data_loader):
-                if num_bn_forget_steps is not None and i >= num_bn_forget_steps:
-                    break
-                args_kwargs_tuple = data_loader.get_inputs(loaded_item)
-                self._infer_batch(args_kwargs_tuple, device)
+    def _run_model_inference(self, data_loader, num_init_steps, device):
+        num_bn_forget_steps = self.num_bn_forget_steps
 
-            self.model.apply(self._apply_to_batchnorms(restore_original_bn_momenta))
+        with self._bn_training_state_switcher():
+            if num_bn_forget_steps is not None and num_bn_forget_steps > 0:
+                with self._bn_momentum_switcher():
+                    for i, loaded_item in enumerate(data_loader):
+                        if i >= num_bn_forget_steps:
+                            break
+                        args_kwargs_tuple = data_loader.get_inputs(loaded_item)
+                        self._infer_batch(args_kwargs_tuple, device)
 
             for i, loaded_item in ProgressBar(
                     enumerate(data_loader),

nncf/model_creation.py

@@ -14,6 +14,7 @@
 from typing import Callable, Any, Tuple, Dict
 
 from torch.nn import Module
+from torch.distributed import barrier
 
 from nncf.checkpoint_loading import load_state
 from nncf.composite_compression import PTCompositeCompressionAlgorithmBuilder
@@ -24,6 +25,7 @@
 from nncf.graph.graph_builder import GraphBuilder
 from nncf.nncf_network import NNCFNetwork
 from nncf.utils import is_main_process
+from nncf.utils import is_dist_avail_and_initialized
 from nncf.algo_selector import COMPRESSION_ALGORITHMS
 
 from nncf.common.utils.logger import logger
@@ -141,4 +143,19 @@ def create_compressed_model(model: Module, config: NNCFConfig,
             graph = compressed_graph_builder.build_graph(compressed_model, compressed_model.get_tracing_context())
             graph.visualize_graph(osp.join(config.get("log_dir", "."), "compressed_graph.dot"))
 
+    # Synchronize all processes if run in distributed mode
+    if is_dist_avail_and_initialized():
+        try:
+            barrier()
+        # Exception can be raised during running barrier
+        # if the backend not in the supported list https://pytorch.org/docs/stable/distributed.html
+        except RuntimeError as err:
+            logger.warning(err)
+            logger.warning(
+                "NNCF continues work, while does not guarantee that "
+                "the processes will finish model's compression at the same time. "
+                "If your training pipeline demands the processes be synchronized, please, "
+                "keep attention to that error")
+            return compression_ctrl, compressed_model
+
     return compression_ctrl, compressed_model

nncf/nncf_network.py

@@ -504,10 +504,15 @@ def _set_nncf_wrapped_model(self, value):
     def get_clean_shallow_copy(self) -> 'NNCFNetwork':
         # WARNING: Will reset pre- and post-ops of the underlying model. Use save_nncf_module_additions
         # and load_nncf_module_additions to preserve these, or temporary_clean_view().
-        return NNCFNetwork(self.get_nncf_wrapped_model(), self.input_infos,
-                           self._user_dummy_forward_fn, self._wrap_inputs_fn,
-                           self.scopes_without_shape_matching, self.ignored_scopes, self.target_scopes,
-                           reset=True)
+        from nncf.utils import save_module_training_state, load_module_training_state
+        saved_state = {}
+        save_module_training_state(self, saved_state)
+        model_copy = NNCFNetwork(self.get_nncf_wrapped_model(), self.input_infos,
+                    self._user_dummy_forward_fn, self._wrap_inputs_fn,
+                    self.scopes_without_shape_matching, self.ignored_scopes, self.target_scopes,
+                    reset=True)
+        load_module_training_state(model_copy, saved_state)
+        return model_copy
 
     def get_modules_in_nncf_modules_by_type(self, types) -> Dict['Scope', nn.Module]:
         nncf_modules = self.get_nncf_modules()

nncf/quantization/algo.py

@@ -606,7 +606,7 @@ def _get_transformation_layout(self, target_model: NNCFNetwork) -> PTTransformat
         target_model.register_compression_module_type(ExtraCompressionModuleType.EXTERNAL_QUANTIZER)
         single_config_quantizer_setup = self._get_quantizer_setup(target_model)
         minmax_values_for_range_init = {}
-        if self.should_init:
+        if is_main_process() and self.should_init:
             stats_for_range_init = self._get_statistics_for_final_range_init(target_model,
                                                                              single_config_quantizer_setup,
                                                                              self._range_init_params)
@@ -1365,6 +1365,9 @@ def __init__(self):
         self.dump_dir = Path(DEBUG_LOG_DIR) / Path("debug_dumps")
         self.dump_dir.mkdir(parents=True, exist_ok=True)
         self.scale_dump_dir = self.dump_dir / Path("scale")
+        if self.scale_dump_dir.exists():
+            shutil.rmtree(str(self.scale_dump_dir))
+        self.scale_dump_dir.mkdir(parents=True, exist_ok=True)
         self.prop_graph_dump_dir = self.dump_dir / Path("quant_prop")
         if self.prop_graph_dump_dir.exists():
             shutil.rmtree(str(self.prop_graph_dump_dir))
@@ -1383,9 +1386,6 @@ def init_actual(self, owner_model: NNCFNetwork):
             nncf_module_quantizations_id_list)
         self.call_trackers[self.ACTIVATION_QUANTIZERS_TRACKER_NAME].init_with_key_list(
             activation_quantizer_id_list)
-        if self.scale_dump_dir.exists():
-            shutil.rmtree(str(self.scale_dump_dir))
-        self.scale_dump_dir.mkdir(parents=True, exist_ok=True)
         self._strict_forward = True
 
     def pre_forward_actions(self, module: 'NNCFNetwork'):
@@ -1428,7 +1428,7 @@ def dump_scale(self, quantizer_scale_params: Dict[str, torch.Tensor], quantizer_
         quantizer_normalized_name = re.sub(r'[^\w\-_\. ]', '_', quantizer_name)
         for scale_param_name, scale_param in quantizer_scale_params.items():
             fname = "{}_{}.txt".format(quantizer_normalized_name, scale_param_name)
-            with safe_open(self.scale_dump_dir / fname, "ba") as file:
+            with safe_open(self.scale_dump_dir / fname, "ab") as file:
                 np.savetxt(file, scale_param.cpu().numpy().flatten())
 
     def reset_counters(self):