Move PyTorch API to torch dir

MANIFEST.in

@@ -1,3 +1,3 @@
-graft nncf/extensions
+graft nncf/torch/extensions
 graft nncf/common/hardware/config
 include LICENSE

examples/classification/common.py

@@ -1,7 +1,7 @@
 import torch.backends.cudnn as cudnn
 
 from examples.common.model_loader import load_resuming_model_state_dict_and_checkpoint_from_path
-from nncf.utils import manual_seed
+from nncf.torch.utils import manual_seed
 
 
 def set_seed(config):

examples/classification/main.py

@@ -48,9 +48,9 @@
 from examples.common.utils import write_metrics
 from nncf import create_compressed_model
 from nncf.api.compression import CompressionStage
-from nncf.dynamic_graph.graph_tracer import create_input_infos
-from nncf.initialization import register_default_init_args, default_criterion_fn
-from nncf.utils import safe_thread_call, is_main_process
+from nncf.torch.dynamic_graph.graph_tracer import create_input_infos
+from nncf.torch.initialization import register_default_init_args, default_criterion_fn
+from nncf.torch.utils import safe_thread_call, is_main_process
 from examples.classification.common import set_seed, load_resuming_checkpoint
 
 model_names = sorted(name for name in models.__dict__

examples/classification/staged_quantization_worker.py

@@ -31,12 +31,12 @@
 from examples.common.model_loader import load_model
 from examples.common.utils import configure_logging, print_args, make_additional_checkpoints, get_name, \
     print_statistics, is_pretrained_model_requested, log_common_mlflow_params, SafeMLFLow, configure_device
-from nncf.binarization.algo import BinarizationController
+from nncf.torch.binarization.algo import BinarizationController
 from nncf.api.compression import CompressionStage
-from nncf.initialization import register_default_init_args, default_criterion_fn
-from nncf.model_creation import create_compressed_model
-from nncf.quantization.algo import QuantizationController
-from nncf.utils import is_main_process
+from nncf.torch.initialization import register_default_init_args, default_criterion_fn
+from nncf.torch.model_creation import create_compressed_model
+from nncf.torch.quantization.algo import QuantizationController
+from nncf.torch.utils import is_main_process
 from examples.classification.common import set_seed, load_resuming_checkpoint
 
 

examples/common/model_loader.py

@@ -20,7 +20,7 @@
 from examples.common.example_logger import logger
 import examples.common.restricted_pickle_module as restricted_pickle_module
 from nncf import load_state
-from nncf.utils import safe_thread_call
+from nncf.torch.utils import safe_thread_call
 from examples.classification.models.mobilenet_v2_32x32 import MobileNetV2For32x32
 
 

examples/common/models/classification/mobilenetv3.py

@@ -24,7 +24,7 @@
 import torch.nn as nn
 import math
 
-from nncf.dynamic_graph.patch_pytorch import register_operator
+from nncf.torch.dynamic_graph.patch_pytorch import register_operator
 
 __all__ = ['mobilenetv3_Large', 'mobilenetv3']
 

examples/common/models/segmentation/icnet.py

@@ -27,7 +27,7 @@
 import torch.nn.functional as F
 
 from examples.common.example_logger import logger
-from nncf.utils import is_tracing_state
+from nncf.torch.utils import is_tracing_state
 
 
 class ConvBN(nn.Module):

examples/common/models/segmentation/unet.py

@@ -20,7 +20,7 @@
 import torch
 import torch.nn.functional as F
 
-from nncf.utils import is_tracing_state
+from nncf.torch.utils import is_tracing_state
 
 from examples.common.example_logger import logger
 

examples/common/utils.py

@@ -36,7 +36,7 @@
 import torch
 
 from examples.common.example_logger import logger as default_logger
-from nncf.utils import is_main_process
+from nncf.torch.utils import is_main_process
 
 # pylint: disable=import-error
 from returns.maybe import Maybe, Nothing

examples/object_detection/dataset.py

@@ -19,7 +19,7 @@
 from examples.object_detection.datasets.coco import COCODataset
 from examples.object_detection.datasets.voc0712 import VOCDetection, VOCAnnotationTransform
 from examples.object_detection.utils.augmentations import SSDAugmentation
-from nncf.dynamic_graph.graph_tracer import create_input_infos
+from nncf.torch.dynamic_graph.graph_tracer import create_input_infos
 
 VOC_MEAN = (0.406, 0.456, 0.485)
 VOC_STD = (0.255, 0.224, 0.229)

examples/object_detection/eval.py

@@ -24,7 +24,7 @@
 from torch.nn import functional as F
 
 from examples.common.example_logger import logger
-from nncf.utils import is_main_process, is_dist_avail_and_initialized, get_world_size
+from nncf.torch.utils import is_main_process, is_dist_avail_and_initialized, get_world_size
 
 
 def str2bool(v):

examples/object_detection/layers/extensions/__init__.py

@@ -4,7 +4,7 @@
 
 from torch.utils.cpp_extension import load
 
-from nncf.extensions import CudaNotAvailableStub
+from nncf.torch.extensions import CudaNotAvailableStub
 
 ext_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
 if torch.cuda.is_available():

examples/object_detection/layers/functions/detection.py

@@ -14,8 +14,8 @@
 import torch
 from torch import nn
 
-from nncf.utils import no_jit_trace, add_domain
-from nncf.dynamic_graph.context import no_nncf_trace
+from nncf.torch.utils import no_jit_trace, add_domain
+from nncf.torch.dynamic_graph.context import no_nncf_trace
 
 from ..box_utils import decode, nms
 

examples/object_detection/layers/functions/prior_box.py

@@ -19,7 +19,7 @@
 import torch
 from torch import nn
 
-from nncf.utils import add_domain
+from nncf.torch.utils import add_domain
 
 class PriorBox(nn.Module):
     def __init__(self, min_size, max_size, aspect_ratio, flip, clip, variance, step, offset,

examples/object_detection/layers/modules/l2norm.py

@@ -15,7 +15,7 @@
 import torch.nn as nn
 import torch.nn.init as init
 
-from nncf.utils import add_domain
+from nncf.torch.utils import add_domain
 from nncf import register_module
 
 

examples/object_detection/layers/modules/ssd_head.py

@@ -16,7 +16,7 @@
 from torch.nn import functional as F
 
 from examples.object_detection.layers import DetectionOutput, PriorBox
-from nncf.dynamic_graph.context import no_nncf_trace
+from nncf.torch.dynamic_graph.context import no_nncf_trace
 
 
 class SSDDetectionOutput(nn.Module):

examples/object_detection/main.py

@@ -31,7 +31,7 @@
 from examples.common.execution import get_execution_mode
 from examples.common.execution import prepare_model_for_execution, start_worker
 from nncf.api.compression import CompressionStage
-from nncf.initialization import register_default_init_args
+from nncf.torch.initialization import register_default_init_args
 from examples.common.optimizer import get_parameter_groups, make_optimizer
 from examples.common.utils import get_name, make_additional_checkpoints, print_statistics, configure_paths, \
     create_code_snapshot, is_on_first_rank, configure_logging, print_args, is_pretrained_model_requested, \
@@ -42,8 +42,8 @@
 from examples.object_detection.layers.modules import MultiBoxLoss
 from examples.object_detection.model import build_ssd
 from nncf import create_compressed_model, load_state
-from nncf.dynamic_graph.graph_tracer import create_input_infos
-from nncf.utils import is_main_process
+from nncf.torch.dynamic_graph.graph_tracer import create_input_infos
+from nncf.torch.utils import is_main_process
 
 
 def str2bool(v):

examples/object_detection/models/ssd_mobilenet.py

@@ -17,7 +17,7 @@
 from examples.common import restricted_pickle_module
 from examples.common.example_logger import logger
 from examples.object_detection.layers.modules.ssd_head import MultiOutputSequential, SSDDetectionOutput
-from nncf.checkpoint_loading import load_state
+from nncf.torch.checkpoint_loading import load_state
 
 
 def conv_bn(inp, oup, kernel, stride, padding):

examples/object_detection/models/ssd_vgg.py

@@ -19,7 +19,7 @@
 from examples.common.example_logger import logger
 from examples.object_detection.layers import L2Norm
 from examples.object_detection.layers.modules.ssd_head import MultiOutputSequential, SSDDetectionOutput
-from nncf.checkpoint_loading import load_state
+from nncf.torch.checkpoint_loading import load_state
 
 BASE_NUM_OUTPUTS = {
     300: [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M', 512, 512, 512],

examples/semantic_segmentation/main.py

@@ -35,7 +35,7 @@
 from examples.common.execution import get_execution_mode, \
     prepare_model_for_execution, start_worker
 from nncf.api.compression import CompressionStage
-from nncf.initialization import register_default_init_args
+from nncf.torch.initialization import register_default_init_args
 from examples.common.model_loader import load_model, load_resuming_model_state_dict_and_checkpoint_from_path
 from examples.common.optimizer import make_optimizer
 from examples.common.utils import configure_logging, configure_paths, make_additional_checkpoints, print_args, \
@@ -46,7 +46,7 @@
 from examples.semantic_segmentation.train import Train
 from examples.semantic_segmentation.utils.checkpoint import save_checkpoint
 from nncf import create_compressed_model
-from nncf.utils import is_main_process
+from nncf.torch.utils import is_main_process
 
 
 def get_arguments_parser():

examples/semantic_segmentation/utils/checkpoint.py

@@ -15,7 +15,7 @@
 import torch
 
 from examples.common import restricted_pickle_module
-from nncf.checkpoint_loading import load_state
+from nncf.torch.checkpoint_loading import load_state
 
 
 def save_checkpoint(model, optimizer, epoch, miou, compression_stage, compression_scheduler, config):

nncf/__init__.py

@@ -11,13 +11,13 @@
  limitations under the License.
 """
 
-from .common.utils.backend import __nncf_backend__
-from .config import NNCFConfig
-from .version import __version__
+from nncf.common.utils.backend import __nncf_backend__
+from nncf.config import NNCFConfig
+from nncf.version import __version__
 
 if __nncf_backend__ == 'Torch':
 
-    from .version import BKC_TORCH_VERSION
+    from nncf.version import BKC_TORCH_VERSION
     import torch
     from pkg_resources import parse_version
     if parse_version(BKC_TORCH_VERSION).base_version != parse_version(torch.__version__).base_version:
@@ -31,38 +31,38 @@
     # NNCF builds extensions based on torch load() function
     # This function has a bug inside which patch_extension_build_function() solves
     # This bug will be fixed in torch 1.8.0
-    from .dynamic_graph.patch_pytorch import patch_extension_build_function
+    from nncf.torch.dynamic_graph.patch_pytorch import patch_extension_build_function
     # It must be called before importing packages containing CUDA extensions
     patch_extension_build_function()
 
     # Required for correct COMPRESSION_ALGORITHMS registry functioning
-    from .binarization import algo as binarization_algo
-    from .quantization import algo as quantization_algo
-    from .sparsity.const import algo as const_sparsity_algo
-    from .sparsity.magnitude import algo as magnitude_sparsity_algo
-    from .sparsity.rb import algo as rb_sparsity_algo
-    from .pruning.filter_pruning import algo as filter_pruning_algo
+    from nncf.torch.binarization import algo as binarization_algo
+    from nncf.torch.quantization import algo as quantization_algo
+    from nncf.torch.sparsity.const import algo as const_sparsity_algo
+    from nncf.torch.sparsity.magnitude import algo as magnitude_sparsity_algo
+    from nncf.torch.sparsity.rb import algo as rb_sparsity_algo
+    from nncf.torch.pruning.filter_pruning import algo as filter_pruning_algo
 
     # Functions most commonly used in integrating NNCF into training pipelines are
     # listed below for importing convenience
 
-    from .model_creation import create_compressed_model
-    from .checkpoint_loading import load_state
-    from .common.utils.logger import disable_logging
-    from .common.utils.logger import set_log_level
-    from .initialization import register_default_init_args
-    from .layers import register_module
-    from .dynamic_graph.patch_pytorch import register_operator
-    from .dynamic_graph.io_handling import nncf_model_input
-    from .dynamic_graph.io_handling import nncf_model_output
-    from .dynamic_graph.context import no_nncf_trace
-    from .dynamic_graph.context import forward_nncf_trace
+    from nncf.torch.model_creation import create_compressed_model
+    from nncf.torch.checkpoint_loading import load_state
+    from nncf.common.utils.logger import disable_logging
+    from nncf.common.utils.logger import set_log_level
+    from nncf.torch.initialization import register_default_init_args
+    from nncf.torch.layers import register_module
+    from nncf.torch.dynamic_graph.patch_pytorch import register_operator
+    from nncf.torch.dynamic_graph.io_handling import nncf_model_input
+    from nncf.torch.dynamic_graph.io_handling import nncf_model_output
+    from nncf.torch.dynamic_graph.context import no_nncf_trace
+    from nncf.torch.dynamic_graph.context import forward_nncf_trace
 
     # NNCF relies on tracing PyTorch operations. Each code that uses NNCF
     # should be executed with PyTorch operators wrapped via a call to "patch_torch_operators",
     # so this call is moved to package __init__ to ensure this.
-    from .dynamic_graph.patch_pytorch import patch_torch_operators
+    from nncf.torch.dynamic_graph.patch_pytorch import patch_torch_operators
 
-    from .extensions import force_build_cpu_extensions, force_build_cuda_extensions
+    from nncf.torch.extensions import force_build_cpu_extensions, force_build_cuda_extensions
 
     patch_torch_operators()

nncf/common/factory.py

@@ -17,7 +17,7 @@
 from nncf.common.exporter import Exporter
 from nncf.common.utils.backend import __nncf_backend__
 if __nncf_backend__ == 'Torch':
-    from nncf.exporter import PTExporter
+    from nncf.torch.exporter import PTExporter
 elif __nncf_backend__ == 'TensorFlow':
     from beta.nncf.tensorflow.exporter import TFExporter
 

nncf/common/graph/operator_metatypes.py

@@ -113,7 +113,7 @@ def get_operator_metatypes() -> Optional[OperatorMetatypeRegistry]:
     :return: The operator metatype registry.
     """
     if __nncf_backend__ == 'Torch':
-        from nncf.graph.operator_metatypes \
+        from nncf.torch.graph.operator_metatypes \
             import PT_OPERATOR_METATYPES
         return PT_OPERATOR_METATYPES
     if __nncf_backend__ == 'TensorFlow':

nncf/common/graph/version_agnostic_op_names.py

@@ -29,7 +29,7 @@ def get_version_agnostic_name(version_specific_name: str):
     """
     def get_func_impl():
         if __nncf_backend__ == 'Torch':
-            from nncf.graph.version_agnostic_op_names \
+            from nncf.torch.graph.version_agnostic_op_names \
                 import get_version_agnostic_name as torch_fn_impl
             return torch_fn_impl
 

nncf/common/pruning/pruning_node_selector.py

@@ -30,10 +30,10 @@
 from nncf.common.utils.backend import __nncf_backend__
 
 if __nncf_backend__ == 'Torch':
-    from nncf.pruning.utils import is_depthwise_conv
-    from nncf.pruning.utils import is_conv_with_downsampling
-    from nncf.graph.graph import get_module_identifier
-    from nncf.utils import should_consider_scope
+    from nncf.torch.pruning.utils import is_depthwise_conv
+    from nncf.torch.pruning.utils import is_conv_with_downsampling
+    from nncf.torch.graph.graph import get_module_identifier
+    from nncf.torch.utils import should_consider_scope
 elif __nncf_backend__ == 'TensorFlow':
     from beta.nncf.tensorflow.pruning.utils import is_depthwise_conv
     from beta.nncf.tensorflow.pruning.utils import is_conv_with_downsampling

nncf/torch/__init__.py

@@ -0,0 +1,12 @@
+"""
+ Copyright (c) 2021 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.
+"""

nncf/algo_selector.py → nncf/torch/algo_selector.py

@@ -14,15 +14,15 @@
 # pylint:disable=relative-beyond-top-level
 import torch
 
-from nncf.graph.transformations.layout import PTTransformationLayout
-from nncf.nncf_network import NNCFNetwork
+from nncf.torch.graph.transformations.layout import PTTransformationLayout
+from nncf.torch.nncf_network import NNCFNetwork
 
 from nncf.api.compression import CompressionStage
 from nncf.api.compression import CompressionScheduler
-from nncf.compression_method_api import PTCompressionAlgorithmBuilder
-from nncf.compression_method_api import PTCompressionAlgorithmController
+from nncf.torch.compression_method_api import PTCompressionAlgorithmBuilder
+from nncf.torch.compression_method_api import PTCompressionAlgorithmController
 
-from nncf.compression_method_api import PTCompressionLoss
+from nncf.torch.compression_method_api import PTCompressionLoss
 from nncf.common.schedulers import StubCompressionScheduler
 from nncf.common.utils.registry import Registry
 

nncf/automl/agent/ddpg/ddpg.py → nncf/torch/automl/agent/ddpg/ddpg.py

@@ -20,7 +20,7 @@
 from torch.optim import Adam
 from types import SimpleNamespace
 
-from nncf.automl.agent.ddpg.memory import SequentialMemory
+from nncf.torch.automl.agent.ddpg.memory import SequentialMemory
 
 
 criterion = nn.MSELoss()