Apply comments

examples/experimental/torch/classification/bootstrap_nas.py

@@ -46,7 +46,7 @@
 from nncf.experimental.torch.nas.bootstrapNAS import SearchAlgorithm
 from nncf.torch.initialization import default_criterion_fn
 from nncf.torch.initialization import wrap_dataloader_for_init
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from nncf.torch.utils import is_main_process
 
 
@@ -174,7 +174,7 @@ def validate_model_fn_top1(model_, loader_):
         top1, _, _ = validate_model_fn(model_, loader_)
         return top1
 
-    nncf_network = create_nncf_network_with_inputs_from_config(model, nncf_config)
+    nncf_network = create_nncf_network(model, nncf_config)
 
     resuming_checkpoint_path = config.resuming_checkpoint_path
     if resuming_checkpoint_path is None:

examples/experimental/torch/classification/bootstrap_nas_search.py

@@ -42,7 +42,7 @@
 from nncf.experimental.torch.nas.bootstrapNAS.training.model_creator_helpers import resume_compression_from_state
 from nncf.torch.checkpoint_loading import load_state
 from nncf.torch.initialization import wrap_dataloader_for_init
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from nncf.torch.utils import is_main_process
 
 
@@ -140,7 +140,7 @@ def validate_model_fn_top1(model_, loader_):
         top1, _, _ = validate_model_fn(model_, loader_)
         return top1
 
-    nncf_network = create_nncf_network_with_inputs_from_config(model, nncf_config)
+    nncf_network = create_nncf_network(model, nncf_config)
 
     if config.search_mode_active:
         compression_state = torch.load(config.search_elasticity_state_path)

nncf/experimental/torch/nas/bootstrapNAS/search/supernet.py

@@ -18,7 +18,7 @@
 from nncf.experimental.torch.nas.bootstrapNAS.elasticity.multi_elasticity_handler import SubnetConfig
 from nncf.experimental.torch.nas.bootstrapNAS.training.model_creator_helpers import resume_compression_from_state
 from nncf.torch.checkpoint_loading import load_state
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from nncf.torch.nncf_network import NNCFNetwork
 
 TModel = TypeVar("TModel")
@@ -59,7 +59,7 @@ def from_checkpoint(
         :param supernet_weights_path: trained weights to resume the super-network.
         :return: SuperNetwork with wrapped functionality.
         """
-        nncf_network = create_nncf_network_with_inputs_from_config(model, nncf_config)
+        nncf_network = create_nncf_network(model, nncf_config)
         compression_state = torch.load(supernet_elasticity_path, map_location=torch.device(nncf_config.device))
         model, elasticity_ctrl = resume_compression_from_state(nncf_network, compression_state)
         model_weights = torch.load(supernet_weights_path, map_location=torch.device(nncf_config.device))

nncf/torch/dynamic_graph/io_handling.py

@@ -24,7 +24,6 @@
 from nncf.common.logging import nncf_logger
 from nncf.common.utils.api_marker import api
 from nncf.config.structures import BNAdaptationInitArgs
-from nncf.config.structures import NNCFExtraConfigStruct
 from nncf.config.structures import QuantizationRangeInitArgs
 from nncf.torch.dynamic_graph.context import forward_nncf_trace
 from nncf.torch.dynamic_graph.patch_pytorch import register_operator
@@ -85,29 +84,16 @@ def get_forward_inputs(self, device: str = None) -> Tuple[Tuple, Dict]:
         pass
 
 
-class ModelInputElement(abc.ABC):
-    def __init__(self, keyword=None):
-        self.keyword = keyword
-
-    @abc.abstractmethod
-    def get_tensor_for_input(self) -> torch.Tensor:
-        pass
-
-    @abc.abstractmethod
-    def is_integer_input(self):
-        pass
-
-
-class FillerInputElement(ModelInputElement):
+class FillerInputElement:
     FILLER_TYPE_ONES = "ones"
     FILLER_TYPE_ZEROS = "zeros"
     FILLER_TYPE_RANDOM = "random"
     FILLER_TYPES = [FILLER_TYPE_ONES, FILLER_TYPE_ZEROS, FILLER_TYPE_RANDOM]
 
     def __init__(self, shape: List[int], type_str: str = "float", keyword=None, filler=None):
-        super().__init__(keyword=keyword)
         self.shape = shape
         self.type = self._string_to_torch_type(type_str)
+        self.keyword = keyword
         if filler is None:
             self.filler = self.FILLER_TYPE_ONES
         else:
@@ -131,7 +117,7 @@ def is_integer_input(self):
         return self.type != torch.float32
 
     def __eq__(self, other: "FillerInputElement"):
-        return self.type == other.type and self.keyword == other.keyword
+        return self.type == other.type and self.keyword == other.keyword and self.filler == other.filler
 
     def get_tensor_for_input(self) -> torch.Tensor:
         if self.filler == FillerInputElement.FILLER_TYPE_ZEROS:
@@ -278,4 +264,4 @@ def data_loader(self) -> NNCFDataLoader:
         pass
 
 
-EXTRA_STRUCTS_WITH_DATALOADERS: List[Type[NNCFExtraConfigStruct]] = [QuantizationRangeInitArgs, BNAdaptationInitArgs]
+EXTRA_STRUCTS_WITH_DATALOADERS: List[Type[HasDataloader]] = [QuantizationRangeInitArgs, BNAdaptationInitArgs]

nncf/torch/exporter.py

@@ -130,7 +130,7 @@ def _export_to_onnx(self, save_path: str, opset_version: int) -> None:
 
         if self._model_args is not None:
             args = tuple(list(args) + list(self._model_args[:-1]))
-            kwargs.extend(**self._model_args[-1])
+            kwargs.update(**self._model_args[-1])
 
         def to_single_batch_tensors(obj: torch.Tensor):
             return obj[0:1]

nncf/torch/model_creation.py

@@ -184,10 +184,6 @@ def get_input_info_from_config(config: NNCFConfig) -> ModelInputInfo:
     )
 
 
-def create_nncf_network_with_inputs_from_config(model: torch.nn.Module, config: NNCFConfig):
-    return create_nncf_network(model, config=config)
-
-
 def create_nncf_network(
     model: torch.nn.Module,
     config: NNCFConfig,

nncf/torch/quantization/quantize_model.py

@@ -35,7 +35,7 @@
 
 
 # TODO(alexsu52): It is a workaround and should be removed.
-class CalibrationDataLoaderAdapter(PTInitializingDataLoader):
+class CalibrationDataLoader(PTInitializingDataLoader):
     """
     An adapter from the nncf.Dataset to the PTInitializingDataLoader.
     """
@@ -60,7 +60,7 @@ def __iter__(self):
     def __len__(self):
         if self._length is None:
             data = self._dataset.get_inference_data()
-            self._length = CalibrationDataLoaderAdapter._get_length(data)
+            self._length = CalibrationDataLoader._get_length(data)
         return self._length
 
     def get_inputs(self, dataloader_output: Any) -> Tuple[Tuple, Dict]:
@@ -220,7 +220,7 @@ def quantize_impl(
         preset, target_device, subset_size, model_type, ignored_scope, advanced_parameters
     )
 
-    calibration_data_loader = CalibrationDataLoaderAdapter(calibration_dataset)
+    calibration_data_loader = CalibrationDataLoader(calibration_dataset)
     nncf_config.register_extra_structs(
         [
             QuantizationRangeInitArgs(data_loader=calibration_data_loader),

tests/torch/nas/creators.py

@@ -28,7 +28,7 @@
 from nncf.torch.dynamic_graph.graph_tracer import create_dummy_forward_fn
 from nncf.torch.dynamic_graph.io_handling import FillerInputInfo
 from nncf.torch.graph.transformations.layout import PTTransformationLayout
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from nncf.torch.model_transformer import PTModelTransformer
 from nncf.torch.nncf_network import NNCFNetwork
 from tests.torch import test_models
@@ -49,7 +49,7 @@ def create_bootstrap_training_model_and_ctrl(
     model, nncf_config: NNCFConfig, register_bn_adapt: bool = True
 ) -> Tuple[NNCFNetwork, BNASTrainingController]:
     algo_name = nncf_config.get("bootstrapNAS", {}).get("training", {}).get("algorithm", "progressive_shrinking")
-    nncf_network = create_nncf_network_with_inputs_from_config(model, nncf_config)
+    nncf_network = create_nncf_network(model, nncf_config)
     if register_bn_adapt:
         register_bn_adaptation_init_args(nncf_config)
     ctrl, model = create_compressed_model_from_algo_names(nncf_network, nncf_config, [algo_name], False)

tests/torch/nas/descriptors.py

@@ -14,7 +14,7 @@
 from nncf.experimental.torch.nas.bootstrapNAS.elasticity.base_handler import SingleElasticityBuilder
 from nncf.experimental.torch.nas.bootstrapNAS.elasticity.base_handler import create_elasticity_builder_from_config
 from nncf.experimental.torch.nas.bootstrapNAS.elasticity.elasticity_dim import ElasticityDim
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from tests.torch.helpers import get_empty_config
 from tests.torch.nas.helpers import move_model_to_cuda_if_available
 from tests.torch.nas.models.synthetic import ThreeConvModel
@@ -56,7 +56,7 @@ def build_handler(self):
         if not input_size:
             input_size = model.INPUT_SIZE
         config = get_empty_config(input_sample_sizes=input_size)
-        nncf_network = create_nncf_network_with_inputs_from_config(model, config)
+        nncf_network = create_nncf_network(model, config)
         builder = self.create_builder()
         handler = builder.build(nncf_network)
         return handler, builder

tests/torch/nas/test_all_elasticity.py

@@ -19,7 +19,7 @@
 from nncf.experimental.torch.nas.bootstrapNAS.elasticity.visualization import SubnetGraph
 from nncf.experimental.torch.nas.bootstrapNAS.training.model_creator_helpers import resume_compression_from_state
 from nncf.torch.exporter import PTExporter
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from tests.shared.nx_graph import compare_nx_graph_with_reference
 from tests.torch.helpers import register_bn_adaptation_init_args
 from tests.torch.nas.creators import create_bnas_model_and_ctrl_by_test_desc
@@ -387,7 +387,7 @@ def test_can_restore_from_state():
     empty_nncf_config = NNCFConfig({"input_info": {"sample_size": THREE_CONV_TEST_DESC.input_sizes}})
     register_bn_adaptation_init_args(empty_nncf_config)
     clean_model = THREE_CONV_TEST_DESC.model_creator()
-    nncf_network = create_nncf_network_with_inputs_from_config(clean_model, empty_nncf_config)
+    nncf_network = create_nncf_network(clean_model, empty_nncf_config)
     _, training_ctrl = resume_compression_from_state(nncf_network, old_state, empty_nncf_config)
 
     new_state = training_ctrl.get_compression_state()
@@ -418,7 +418,7 @@ def test_can_restore_and_get_the_same_output():
     clean_model = THREE_CONV_TEST_DESC.model_creator()
     if torch.cuda.is_available():
         clean_model.cuda()
-    nncf_network = create_nncf_network_with_inputs_from_config(clean_model, empty_nncf_config)
+    nncf_network = create_nncf_network(clean_model, empty_nncf_config)
     model, _ = resume_compression_from_state(nncf_network, old_state, empty_nncf_config)
 
     actual_output = model(input_)

tests/torch/nas/test_ps_controller.py

@@ -20,7 +20,7 @@
 from nncf import NNCFConfig
 from nncf.config.structures import BNAdaptationInitArgs
 from nncf.experimental.torch.nas.bootstrapNAS import EpochBasedTrainingAlgorithm
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from nncf.torch.utils import get_model_device
 from tests.torch.helpers import create_ones_mock_dataloader
 from tests.torch.nas.helpers import move_model_to_cuda_if_available
@@ -123,7 +123,7 @@ def test_bn_adapt(self, mocker, bn_adapt_section_is_called, schedule_params):
             "nncf.common.initialization.batchnorm_adaptation.BatchnormAdaptationAlgorithm.run"
         )
         model, _, nncf_config = prepare_test_model(test_desc, bn_adapt_section_is_called)
-        model = create_nncf_network_with_inputs_from_config(model, nncf_config)
+        model = create_nncf_network(model, nncf_config)
 
         training_algorithm = EpochBasedTrainingAlgorithm.from_config(model, nncf_config)
         training_algorithm._training_ctrl.prepare_for_validation()
@@ -139,7 +139,7 @@ def test_knowledge_distillation_training_process(self):
             input_sizes=ThreeConvModel.INPUT_SIZE,
         )
         model, _, nncf_config = prepare_test_model(test_desc, False, True)
-        model = create_nncf_network_with_inputs_from_config(model, nncf_config)
+        model = create_nncf_network(model, nncf_config)
 
         torch.manual_seed(2)
         number_of_iters = 2

tests/torch/nas/test_state.py

@@ -21,7 +21,7 @@
 from nncf.experimental.torch.nas.bootstrapNAS.elasticity.elastic_kernel import EKBuilderStateNames
 from nncf.experimental.torch.nas.bootstrapNAS.elasticity.elastic_width import EWBuilderStateNames
 from nncf.experimental.torch.nas.bootstrapNAS.elasticity.elasticity_dim import ElasticityDim
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from tests.torch.helpers import BasicConvTestModel
 from tests.torch.helpers import get_empty_config
 from tests.torch.nas.creators import build_elastic_model_from_handler
@@ -223,7 +223,7 @@ def test_can_load_handler_state(desc: ElasticityDesc):
     if not input_size:
         input_size = model.INPUT_SIZE
     config = get_empty_config(input_sample_sizes=input_size)
-    old_nncf_network = create_nncf_network_with_inputs_from_config(model, config)
+    old_nncf_network = create_nncf_network(model, config)
     old_builder = desc.create_builder()
     old_handler = old_builder.build(old_nncf_network)
     elastic_model = build_elastic_model_from_handler(old_nncf_network, old_handler)
@@ -232,7 +232,7 @@ def test_can_load_handler_state(desc: ElasticityDesc):
     ref_output = desc.ref_output_fn(original_model, dummy_input)
     assert torch.allclose(old_output, ref_output)
 
-    new_nncf_network = create_nncf_network_with_inputs_from_config(deepcopy(original_model), config)
+    new_nncf_network = create_nncf_network(deepcopy(original_model), config)
     builder_state = old_builder.get_state()
     # no need in config to restore builder state
     new_builder = desc.create_builder_with_config({})

tests/torch/pruning/experimental/test_nodes_grouping.py

@@ -18,7 +18,7 @@
 from torch import nn
 
 from nncf import NNCFConfig
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from tests.shared.nx_graph import compare_nx_graph_with_reference
 from tests.torch.test_compressed_graph import GeneralModelDesc
 from tests.torch.test_compressed_graph import IModelDesc
@@ -470,7 +470,7 @@ def test_groups(desc: GroupTestDesc, mocker, tmp_path):
     model_desc = desc.model_desc
     model = model_desc.get_model()
     config = NNCFConfig({"input_info": model_desc.create_input_info()})
-    nncf_network = create_nncf_network_with_inputs_from_config(model, config)
+    nncf_network = create_nncf_network(model, config)
     pruning_producing_types = ["linear"]
     get_graph_spy = mocker.spy(BlockHierarchy, "_get_graph_for_visualization")
     not_filtered_groups = get_pruning_groups(
@@ -508,7 +508,7 @@ def test_all_groups_valid(desc: GroupTestDesc):
     model_desc = desc.model_desc
     model = model_desc.get_model()
     config = NNCFConfig({"input_info": model_desc.create_input_info()})
-    nncf_network = create_nncf_network_with_inputs_from_config(model, config)
+    nncf_network = create_nncf_network(model, config)
     pruning_producing_types = ["linear"]
     all_groups = get_pruning_groups(
         nncf_network.nncf.get_graph(), PT_EXPERIMENTAL_PRUNING_OPERATOR_METATYPES, pruning_producing_types

tests/torch/ptq/helpers.py

@@ -20,7 +20,7 @@
 from nncf.torch.graph.operator_metatypes import PTModuleConv2dMetatype
 from nncf.torch.graph.operator_metatypes import PTModuleLinearMetatype
 from nncf.torch.graph.operator_metatypes import PTSumMetatype
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from nncf.torch.tensor_statistics.statistics import PTMinMaxTensorStatistic
 from tests.post_training.test_templates.models import NNCFGraphToTest
 from tests.post_training.test_templates.models import NNCFGraphToTestDepthwiseConv
@@ -80,7 +80,7 @@ def get_nncf_network(model: torch.nn.Module, input_shape: Optional[List[int]] =
     input_shape = [1, 3, 32, 32] if input_shape is None else input_shape
     model.eval()
     nncf_config = NNCFConfig({"input_info": {"sample_size": input_shape.copy()}})
-    nncf_network = create_nncf_network_with_inputs_from_config(
+    nncf_network = create_nncf_network(
         model=model,
         config=nncf_config,
     )

tests/torch/ptq/test_calculation_quantizer_params.py

@@ -32,7 +32,7 @@
 from nncf.quantization.fake_quantize import FakeQuantizeParameters
 from nncf.quantization.fake_quantize import calculate_quantizer_parameters
 from nncf.quantization.fake_quantize import get_quantizer_narrow_range
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from nncf.torch.statistics.aggregator import PTStatisticsAggregator
 from nncf.torch.tensor_statistics.statistics import PTMinMaxTensorStatistic
 from tests.post_training.test_templates.test_calculate_quantizer_parameters import TemplateTestFQParams
@@ -316,7 +316,7 @@ def test_quantizer_parameters_export(tmp_path: Path):
     statistics_aggregator = PTStatisticsAggregator(dataset)
 
     nncf_config = NNCFConfig({"input_info": {"sample_size": [1, 3, 32, 32]}})
-    nncf_network = create_nncf_network_with_inputs_from_config(model, nncf_config)
+    nncf_network = create_nncf_network(model, nncf_config)
     statistic_points = min_max_algo.get_statistic_points(nncf_network, nncf_network.nncf.get_graph())
     statistics_aggregator.register_statistic_points(statistic_points)
     statistics_aggregator.collect_statistics(model, nncf_network.nncf.get_graph())

tests/torch/ptq/test_fq_params_calculation.py

@@ -19,7 +19,7 @@
 from nncf.quantization.advanced_parameters import AdvancedQuantizationParameters
 from nncf.quantization.advanced_parameters import OverflowFix
 from nncf.quantization.algorithms.post_training.algorithm import PostTrainingQuantization
-from nncf.torch.model_creation import create_nncf_network_with_inputs_from_config
+from nncf.torch.model_creation import create_nncf_network
 from nncf.torch.nncf_network import NNCFNetwork
 from nncf.torch.quantization.layers import QUANTIZATION_MODULES
 from nncf.torch.utils import get_all_modules_by_type
@@ -58,7 +58,7 @@ def transform_fn(sample):
     post_training_quantization = PostTrainingQuantization(subset_size=1, **quantization_params)
 
     original_model.eval()
-    nncf_network = create_nncf_network_with_inputs_from_config(original_model, config)
+    nncf_network = create_nncf_network(original_model, config)
     quantized_model = post_training_quantization.apply(nncf_network, nncf_network.nncf.get_graph(), dataset=dataset)
     return quantized_model
 

tests/torch/test_onnx_export.py

@@ -132,10 +132,8 @@ def test_can_export_single_batch_bn(tmp_path):
 
 
 def test_can_export_with_model_args(tmp_path):
-    pytest.xfail(
-        "Torch now parses the function signature and sets up default parameters for unprovided "
-        "arguments on its own. Need to rethink and possibly deprecate model_args parameter."
-    )
+    # Torch now parses the function signature and sets up default parameters for unprovided
+    # arguments on its own. Need to rethink and possibly deprecate model_args parameter.
     test_path = tmp_path.joinpath("test.onnx")
     model = MultiParamForwardModel()
     config = get_basic_quantization_config(input_info=[{"sample_size": [1, 1, 1, 1]}, {"sample_size": [1, 1, 1, 1]}])