import torch.nn as nn
from nncf.common.utils.registry import Registry
COMPRESSION_MODULES = Registry('compression modules')
class ProxyModule:
def __init__(self, module):
self._module = module
def __getattr__(self, name):
return getattr(self._module, name)
class _NNCFModuleMixin:
"""Default class for modules that will be optimized by NNCF.
Attributes:
op_func_name Name of corresponding torch function.
target_weight_dim_for_compression Target dimension of weights that will be compressed in some algorithms.
ignored_algorithms List of algorithms that will skip the module.
"""
op_func_name = ""
target_weight_dim_for_compression = 0
ignored_algorithms = []
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
_NNCFModuleMixin.add_mixin_fields(self)
import numpy as np
from scipy.interpolate import interp1d
from nncf.api.compression import CompressionAlgorithmController
from nncf.common.composite_compression import CompositeCompressionAlgorithmController
from nncf.common.utils.logger import logger as nncf_logger
from nncf.common.utils.registry import Registry
from nncf.config.config import NNCFConfig
from nncf.config.extractors import extract_accuracy_aware_training_params
from nncf.common.accuracy_aware_training.runner import TrainingRunner
from nncf.common.accuracy_aware_training.runner import TrainingRunnerCreator
from nncf.common.accuracy_aware_training.runner import EarlyExitTrainingRunnerCreator
from nncf.common.accuracy_aware_training.runner import AdaptiveCompressionLevelTrainingRunnerCreator
ModelType = TypeVar('ModelType')
ADAPTIVE_COMPRESSION_CONTROLLERS = Registry('adaptive_compression_controllers')
class TrainingLoop(ABC):
"""
The training loop object is instantiated by the user, the training process
is launched via the `run` method.
"""
@abstractmethod
def run(self,
model: ModelType,
train_epoch_fn,
validate_fn,
configure_optimizers_fn=None,
dump_checkpoint_fn=None,
tensorboard_writer=None,
"""
Copyright (c) 2020 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 nncf.common.utils.registry import Registry
NNCF_CUSTOM_OBJECTS = Registry('NNCF Custom Objects')
NNCF_QUANTIZATION_OPERATONS = Registry('NNCF Quantization Operations')
def get_nncf_custom_objects():
return NNCF_CUSTOM_OBJECTS.registry_dict
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.torch.compression_method_api import PTCompressionAlgorithmBuilder
from nncf.torch.compression_method_api import PTCompressionAlgorithmController
from nncf.torch.compression_method_api import PTCompressionLoss
from nncf.common.compression import NO_COMPRESSION_ALGORITHM_NAME
from nncf.common.schedulers import StubCompressionScheduler
from nncf.common.utils.registry import Registry
from nncf.common.statistics import NNCFStatistics
PT_COMPRESSION_ALGORITHMS = Registry('compression algorithm',
add_name_as_attr=True)
class ZeroCompressionLoss(PTCompressionLoss):
def __init__(self, device: str):
super().__init__()
self._device = device
def calculate(self) -> torch.Tensor:
return torch.zeros([], device=self._device)
@PT_COMPRESSION_ALGORITHMS.register(NO_COMPRESSION_ALGORITHM_NAME)
class NoCompressionAlgorithmBuilder(PTCompressionAlgorithmBuilder):
def _get_transformation_layout(
self, target_model: NNCFNetwork) -> PTTransformationLayout:
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 bisect import bisect_right
import numpy as np
from beta.nncf.api.compression import CompressionScheduler
from nncf.common.utils.registry import Registry
SPARSITY_SCHEDULERS = Registry("sparsity_schedulers")
class SparsityScheduler(CompressionScheduler):
def __init__(self, sparsity_algo, params: dict = None):
super().__init__()
if params is None:
self._params = dict()
else:
self._params = params
self.algo = sparsity_algo
self.initial_sparsity = self._params.get('sparsity_init', 0)
self.sparsity_target = self._params.get('sparsity_target', 0.5)
self.sparsity_target_epoch = self._params.get('sparsity_target_epoch',
90)
limitations under the License.
"""
import sys
from argparse import ArgumentParser
from collections import OrderedDict
import torch
from nncf.common.utils.registry import Registry
INCEPTION_NAME = 'inception'
RESNET_NAME = 'resnet'
MOBILENET_NAME = 'mobilenet'
KEYS_REPLACERS = Registry("keys_replacers")
@KEYS_REPLACERS.register(INCEPTION_NAME)
def inception_replacer(k):
if 'RELU' in k:
return k.replace('335', '0')
return k
@KEYS_REPLACERS.register(MOBILENET_NAME)
def mobilenet_replacer(k):
keywords = ['hardtanh', 'batch_norm', '__add__']
if any(x in k for x in keywords):
return k.replace('63', '0').replace('62', '0').replace('111', '0')
return k
limitations under the License.
"""
import torch
import torch.nn as nn
from nncf.dynamic_graph.patch_pytorch import register_operator
from nncf.layer_utils import COMPRESSION_MODULES
from nncf.common.utils.registry import Registry
from nncf.utils import get_per_channel_scale_shape
from nncf.binarization.binarize_functions import XNORBinarizeFn, DOREFABinarizeFn
from nncf.binarization.binarize_functions import ActivationBinarizationScaleThresholdFn
from nncf.nncf_logger import logger as nncf_logger
BINARIZATION_MODULES = Registry('binarization_modules')
class BinarizationMode:
XNOR = "xnor"
DOREFA = "dorefa"
class BaseBinarizer(nn.Module):
def __init__(self, enabled=False):
super().__init__()
self.register_buffer('enabled', torch.IntTensor([0]))
if enabled:
self.enable()
def forward(self, x):
NNCF_GENERAL_CONV_MODULES_DICT = dict(NNCF_CONV_MODULES_DICT)
NNCF_GENERAL_CONV_MODULES_DICT.update(NNCF_DECONV_MODULES_DICT)
NNCF_LINEAR_MODULES_DICT = {NNCFLinear: nn.Linear}
NNCF_PRUNING_MODULES_DICT = {
NNCFConv1d: nn.Conv1d,
NNCFConv2d: nn.Conv2d,
NNCFConv3d: nn.Conv3d,
NNCFConvTranspose2d: nn.ConvTranspose2d,
NNCFConvTranspose3d: nn.ConvTranspose3d,
}
NNCF_PRUNING_MODULES_MAP = {k.__name__: v.__name__ for k, v in NNCF_CONV_MODULES_DICT.items()}
NNCF_PRUNING_MODULES = list(NNCF_CONV_MODULES_MAP.keys())
UNWRAPPED_USER_MODULES = Registry('user_modules')
NNCF_WRAPPED_USER_MODULES_DICT = {}
def register_module(*quantizable_field_names: str, ignored_algorithms: list = None):
# quantizable_field_names will work for `weight` attributes only. Should later extend to registering
# customly named attributes if it becomes necessary
def wrap(cls):
UNWRAPPED_USER_MODULES.registry_dict[cls.__name__] = cls
nncf_wrapped_module_class_name = 'NNCFUser{}'.format(cls.__name__)
NNCF_WRAPPED_USER_MODULES_DICT[cls] = type(nncf_wrapped_module_class_name, (_NNCFModuleMixin, cls), {})
get_base_attributes_fn = lambda self: GenericWeightedLayerAttributes(self.weight.requires_grad,
self.weight.shape)
setattr(NNCF_WRAPPED_USER_MODULES_DICT[cls], "get_weight_shape", get_base_attributes_fn)
if ignored_algorithms:
setattr(NNCF_WRAPPED_USER_MODULES_DICT[cls], "ignored_algorithms", ignored_algorithms)
from torch import distributed
from nncf.torch.checkpoint_loading import OPTIONAL_PARAMETERS_REGISTRY
from nncf.common.utils.debug import is_debug
from nncf.torch.functions import clamp
from nncf.common.utils.logger import logger as nncf_logger
from nncf.common.quantization.structs import QuantizationMode, QuantizerConfig, QuantizerSpec
from nncf.common.quantization.quantizers import calculate_symmetric_level_ranges
from nncf.common.quantization.quantizers import calculate_asymmetric_level_ranges
from nncf.torch.quantization.quantize_functions import symmetric_quantize, asymmetric_quantize, \
ExportQuantizeToFakeQuantize, get_scale_zp_from_input_low_input_high, ExportQuantizeToONNXQuantDequant, TuneRange
from nncf.torch.layer_utils import COMPRESSION_MODULES, CompressionParameter
from nncf.common.utils.registry import Registry
from nncf.torch.utils import get_flat_tensor_contents_string, no_jit_trace, is_tracing_state
QUANTIZATION_MODULES = Registry('quantization_modules')
INITIALIZABLE_MODULES = Registry('initializable_modules')
class QuantizerExportMode(Enum):
FAKE_QUANTIZE = "fake_quantize"
ONNX_QUANTIZE_DEQUANTIZE_PAIRS = "quantize_dequantize"
@staticmethod
def from_str(config_value: str) -> 'HWConfigType':
if config_value == QuantizerExportMode.FAKE_QUANTIZE.value:
return QuantizerExportMode.FAKE_QUANTIZE
if config_value == QuantizerExportMode.ONNX_QUANTIZE_DEQUANTIZE_PAIRS.value:
return QuantizerExportMode.ONNX_QUANTIZE_DEQUANTIZE_PAIRS
raise RuntimeError("Unknown quantizer ONNX export mode string")
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 typing import Optional
import numpy as np
import scipy.optimize
from nncf.common.utils.registry import Registry
from nncf.common.schedulers import ExponentialDecaySchedule, BaseCompressionScheduler
PRUNING_SCHEDULERS = Registry('pruning_schedulers')
class PruningScheduler(BaseCompressionScheduler):
"""
This is the class from which all pruning schedulers inherit.
A pruning scheduler is an object which specifies the pruning
level at each training epoch. It involves a scheduling algorithm,
defined in the `_calculate_pruning_level()` method and a state
(some parameters required for current pruning level calculation)
defined in the `__init__()` method.
:param initial_level: Pruning level which already has been
applied to the model. It is the level at which the schedule begins.
:param target_level: Pruning level at which the schedule ends.
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 beta.nncf.api.compression import CompressionAlgorithmBuilder
from beta.nncf.api.compression import CompressionAlgorithmController
from beta.nncf.tensorflow.graph.transformations.layout import TransformationLayout
from beta.nncf.utils.logger import logger
from nncf.common.utils.registry import Registry
TF_COMPRESSION_ALGORITHMS = Registry('compression algorithm')
@TF_COMPRESSION_ALGORITHMS.register('NoCompressionAlgorithm')
class NoCompressionAlgorithmBuilder(CompressionAlgorithmBuilder):
def get_transformation_layout(self, _):
return TransformationLayout()
class NoCompressionAlgorithmController(CompressionAlgorithmController):
pass
def get_compression_algorithm_builder(config):
algorithm_key = config.get('algorithm', 'NoCompressionAlgorithm')
logger.info('Creating compression algorithm: {}'.format(algorithm_key))
return wrap
def get_operator_metatype_by_op_name(
self, op_name: str) -> Type[OperatorMetatype]:
"""
Returns the operator metatype by operator name.
:param op_name: The operator name.
:return: The operator metatype.
"""
if op_name not in self._op_name_to_op_meta_dict:
return UnknownMetatype
return self._op_name_to_op_meta_dict[op_name]
NOOP_METATYPES = Registry('noop_metatypes')
INPUT_NOOP_METATYPES = Registry('input_noop_metatypes')
OUTPUT_NOOP_METATYPES = Registry('output_noop_metatypes')
class UnknownMetatype(OperatorMetatype):
"""
UnknownMetatype is mapped to operations in NNCFGraph, which are unknown for algorithms,
typically these are the operations that haven't been discovered before.
Algorithms should avoid processing graph nodes with this metatype.
"""
name = "unknown"
@classmethod
def get_all_aliases(cls) -> List[str]:
return [cls.name]
import enum
import torch
from abc import ABC, abstractmethod
from nncf.common.utils.registry import Registry
EXTENSIONS = Registry('extensions')
class ExtensionsType(enum.Enum):
CPU = 0
CUDA = 1
class ExtensionLoader(ABC):
@staticmethod
@abstractmethod
def extension_type():
pass
@staticmethod
@abstractmethod
def load():
pass
def _force_build_extensions(ext_type: ExtensionsType):
for class_type in EXTENSIONS.registry_dict.values():
if class_type.extension_type() != ext_type:
continue
class_type.load()
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.
"""
import logging
from nncf.common.utils.registry import Registry
from nncf.common.schedulers import BaseCompressionScheduler
from nncf.api.compression import CompressionStage
logger = logging.getLogger(__name__)
QUANTIZATION_SCHEDULERS = Registry("quantization_schedulers")
@QUANTIZATION_SCHEDULERS.register("staged")
class StagedQuantizationScheduler(BaseCompressionScheduler):
def __init__(self,
quantization_ctrl: 'QuantizationController',
params=None):
super().__init__()
if params is None:
params = {}
self.algo = quantization_ctrl
self.activations_quant_start_epoch = params.get(
'activations_quant_start_epoch', 1)
self.weights_quant_start_epoch = params.get(
'weights_quant_start_epoch', 1)
