Generalization the weight compression algorithm to PyTorch model. (#2400)

cherry pick to release #2333

Author: alexsu52

docs/api/source/conf.py

@@ -137,8 +137,10 @@ def collect_api_entities() -> APIInfo:
     "keras",
     "tensorflow_addons",
     # Need add backend implementation functions to avoid endless loops on registered functions by mock module,
-    "nncf.experimental.tensor.torch_functions",
-    "nncf.experimental.tensor.numpy_functions",
+    "nncf.experimental.tensor.functions.numpy_numeric",
+    "nncf.experimental.tensor.functions.numpy_linalg",
+    "nncf.experimental.tensor.functions.torch_numeric",
+    "nncf.experimental.tensor.functions.torch_linalg",
 ]
 
 with mock(mock_modules):

nncf/common/graph/layer_attributes.py

@@ -271,9 +271,11 @@ class ConvertDtypeLayerAttributes(BaseLayerAttributes):
 
 
 @dataclass
-class ParameterLayerAttributes(BaseLayerAttributes):
+class ConstantLayerAttributes(BaseLayerAttributes):
     """
-    :param name: Parameter name.
+    :param name: Constant name.
+    :param shape: Constant shape.
     """
 
     name: str
+    shape: List[int]

nncf/common/graph/transformations/commands.py

@@ -214,16 +214,3 @@ def __init__(self, command_type: TransformationType, target_point: TargetPoint):
     @property
     def target_point(self) -> TargetPoint:
         return self._target_point
-
-    def check_command_compatibility(self, command: "TransformationCommand") -> bool:
-        return (
-            isinstance(command, TransformationCommand)
-            and self.type == command.type
-            and self.target_point == command.target_point
-        )
-
-    def union(self, other: "TransformationCommand") -> "TransformationCommand":
-        raise NotImplementedError()
-
-    def __add__(self, other: "TransformationCommand") -> "TransformationCommand":
-        return self.union(other)

nncf/common/scopes.py

@@ -10,7 +10,7 @@
 # limitations under the License.
 
 import re
-from typing import List, Optional, Union
+from typing import List, Optional, Set, Union
 
 from nncf.common.graph import NNCFGraph
 from nncf.common.graph import NNCFNode
@@ -21,23 +21,23 @@
 from nncf.scopes import convert_ignored_scope_to_list
 
 
-def matches_any(tested_str: str, str_or_list_to_match_to: Union[List[str], str]) -> bool:
+def matches_any(tested_str: str, strs_to_match_to: Union[List[str], Set[str], str]) -> bool:
     """
-    Return True if tested_str matches at least one element in str_or_list_to_match_to.
+    Return True if tested_str matches at least one element in strs_to_match_to.
 
     :param tested_str: One of the supported entity types to be matched - currently possible to pass either
         NNCFNodeName (to refer to the original model operations) or QuantizerId (to refer to specific quantizers).
-    :param str_or_list_to_match_to: A list of strings specifying for the serializable_id. Entries of the strings
+    :param strs_to_match_to: A list or set of strings specifying for the serializable_id. Entries of the strings
         may be prefixed with `{re}` to enable regex matching.
 
     :return: A boolean value specifying whether a tested_str should matches at least one element
-        in str_or_list_to_match_to.
+        in strs_to_match_to.
     """
 
-    if str_or_list_to_match_to is None:
+    if strs_to_match_to is None:
         return False
 
-    str_list = [str_or_list_to_match_to] if isinstance(str_or_list_to_match_to, str) else str_or_list_to_match_to
+    str_list = [strs_to_match_to] if isinstance(strs_to_match_to, str) else strs_to_match_to
     for item in str_list:
         if "{re}" in item:
             regex = item.replace("{re}", "")
@@ -51,15 +51,15 @@ def matches_any(tested_str: str, str_or_list_to_match_to: Union[List[str], str])
 
 def should_consider_scope(
     serializable_id: Union[QuantizerId, NNCFNodeName],
-    ignored_scopes: List[str],
+    ignored_scopes: Union[List[str], Set[str]],
     target_scopes: Optional[List[str]] = None,
 ) -> bool:
     """
     Used when an entity arising during compression has to be compared to an allowlist or a denylist of strings.
 
     :param serializable_id: One of the supported entity types to be matched - currently possible to pass either
         NNCFNodeName (to refer to the original model operations) or QuantizerId (to refer to specific quantizers)
-    :param ignored_scopes: A list of strings specifying a denylist for the serializable_id. Entries of the list
+    :param ignored_scopes: A list or set of strings specifying a denylist for the serializable_id. Entries of the list
         may be prefixed with `{re}` to enable regex matching.
     :param target_scopes: A list of strings specifying an allowlist for the serializable_id. Entries of the list
         may be prefixed with `{re}` to enable regex matching.

nncf/experimental/common/tensor_statistics/statistical_functions.py

@@ -10,7 +10,7 @@
 # limitations under the License.
 
 from nncf.experimental.tensor import Tensor
-from nncf.experimental.tensor import functions as fns
+from nncf.experimental.tensor.functions import numeric as fns
 
 
 def mean_per_channel(x: Tensor, axis: int) -> Tensor:

nncf/experimental/tensor/README.md

@@ -54,7 +54,7 @@ nncf_tensor.max()  # Tensor(2)
 
 ### Functions over Tensor
 
-All available functions you can found in [functions.py](functions.py).
+All available functions you can found in the functions module.
 
 ```python
 from nncf.experimental.tensor import functions as fns
@@ -105,7 +105,7 @@ tensor_a[0:2]  # Tensor(array([[1],[2]]))
             return fns.foo(self, arg1)
     ```
 
-2. Add function to [function.py](function.py)
+2. Add function to functions module
 
     ```python
     @functools.singledispatch
@@ -133,17 +133,17 @@ tensor_a[0:2]  # Tensor(array([[1],[2]]))
         raise NotImplementedError(f"Function `foo` is not implemented for {type(x)}")
     ```
 
-3. Add backend specific implementation of method to:
+3. Add backend specific implementation of method to correcponding module:
 
-    - [numpy_function.py](numpy_functions.py)
+    - `functions/numpy_*.py`
 
         ```python
         @_register_numpy_types(fns.foo)
         def _(a: TType, arg1: Type) -> np.ndarray:
             return np.foo(a, arg1)
         ```
 
-    - [torch_function.py](torch_functions.py)
+    - `functions/torch_*.py`
 
         ```python
         @fns.foo.register(torch.Tensor)
@@ -155,7 +155,7 @@ tensor_a[0:2]  # Tensor(array([[1],[2]]))
 
 ### Add new backend
 
-1. Add backend specific implementation for all function from [function.py](function.py) in `<NEW_BACKEND>_functions.py` file.
+1. Add backend specific implementation for all function from functions module in `functions/<NEW_BACKEND>_*.py` file.
 
 2. Add `test_tensor.py` in backend-specific t directory for tests that inherited from class `TemplateTestNNCFTensorOperators`
 
@@ -177,8 +177,8 @@ tensor_a[0:2]  # Tensor(array([[1],[2]]))
         "openvino",
         "tensorflow",
         "tensorflow_addons",
-        "nncf.experimental.tensor.torch_functions",
-        "nncf.experimental.tensor.numpy_functions",
-        "nncf.experimental.tensor.<NEW_BACKEND>_functions",
+        "nncf.experimental.tensor.functions.torch_*",
+        "nncf.experimental.tensor.functions.numpy_*",
+        "nncf.experimental.tensor.functions.<NEW_BACKEND>_*",
     ]
     ```

nncf/experimental/tensor/__init__.py

@@ -9,8 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nncf.experimental.tensor.enums import TensorBackendType as TensorBackendType
-from nncf.experimental.tensor.enums import TensorDataType as TensorDataType
-from nncf.experimental.tensor.enums import TensorDeviceType as TensorDeviceType
+from nncf.experimental.tensor.definitions import TensorBackendType as TensorBackendType
+from nncf.experimental.tensor.definitions import TensorDataType as TensorDataType
+from nncf.experimental.tensor.definitions import TensorDeviceType as TensorDeviceType
 from nncf.experimental.tensor.tensor import Tensor as Tensor
 from nncf.experimental.tensor.tensor import unwrap_tensor_data as unwrap_tensor_data

nncf/experimental/tensor/enums.py nncf/experimental/tensor/definitions.py

@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import dataclass
 from enum import Enum
 from enum import auto
 
@@ -31,6 +32,8 @@ class TensorDataType(Enum):
     float32 = auto()
     float64 = auto()
     int8 = auto()
+    int32 = auto()
+    int64 = auto()
     uint8 = auto()
 
 
@@ -41,3 +44,18 @@ class TensorDeviceType(Enum):
 
     CPU = auto()
     GPU = auto()
+
+
+@dataclass
+class TypeInfo:
+    """
+    The class represents the numerical properties of a floating point types.
+
+    :param eps: The smallest representable number such that 1.0 + eps != 1.0.
+    :param max: The largest representable number.
+    :param min: The smallest representable number (typically -max).
+    """
+
+    eps: float
+    max: float
+    min: float

nncf/experimental/tensor/functions/__init__.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2024 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.experimental.tensor.functions import linalg as linalg
+from nncf.experimental.tensor.functions.numeric import abs as abs
+from nncf.experimental.tensor.functions.numeric import all as all
+from nncf.experimental.tensor.functions.numeric import allclose as allclose
+from nncf.experimental.tensor.functions.numeric import any as any
+from nncf.experimental.tensor.functions.numeric import as_tensor_like as as_tensor_like
+from nncf.experimental.tensor.functions.numeric import astype as astype
+from nncf.experimental.tensor.functions.numeric import clip as clip
+from nncf.experimental.tensor.functions.numeric import count_nonzero as count_nonzero
+from nncf.experimental.tensor.functions.numeric import device as device
+from nncf.experimental.tensor.functions.numeric import dtype as dtype
+from nncf.experimental.tensor.functions.numeric import finfo as finfo
+from nncf.experimental.tensor.functions.numeric import flatten as flatten
+from nncf.experimental.tensor.functions.numeric import isclose as isclose
+from nncf.experimental.tensor.functions.numeric import isempty as isempty
+from nncf.experimental.tensor.functions.numeric import item as item
+from nncf.experimental.tensor.functions.numeric import max as max
+from nncf.experimental.tensor.functions.numeric import maximum as maximum
+from nncf.experimental.tensor.functions.numeric import mean as mean
+from nncf.experimental.tensor.functions.numeric import min as min
+from nncf.experimental.tensor.functions.numeric import minimum as minimum
+from nncf.experimental.tensor.functions.numeric import moveaxis as moveaxis
+from nncf.experimental.tensor.functions.numeric import multiply as multiply
+from nncf.experimental.tensor.functions.numeric import ones_like as ones_like
+from nncf.experimental.tensor.functions.numeric import reshape as reshape
+from nncf.experimental.tensor.functions.numeric import round as round
+from nncf.experimental.tensor.functions.numeric import squeeze as squeeze
+from nncf.experimental.tensor.functions.numeric import stack as stack
+from nncf.experimental.tensor.functions.numeric import sum as sum
+from nncf.experimental.tensor.functions.numeric import unstack as unstack
+from nncf.experimental.tensor.functions.numeric import var as var
+from nncf.experimental.tensor.functions.numeric import where as where
+from nncf.experimental.tensor.functions.numeric import zeros_like as zeros_like
+
+
+def _initialize_backends():
+    import contextlib
+
+    import nncf.experimental.tensor.functions.numpy_linalg
+    import nncf.experimental.tensor.functions.numpy_numeric
+
+    with contextlib.suppress(ImportError):
+        import nncf.experimental.tensor.functions.torch_linalg
+        import nncf.experimental.tensor.functions.torch_numeric  # noqa: F401
+
+
+_initialize_backends()

nncf/experimental/tensor/functions/dispatcher.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2024 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.
+import functools
+from typing import List
+
+import numpy as np
+
+from nncf.experimental.tensor import Tensor
+
+
+def tensor_guard(func: callable):
+    """
+    A decorator that ensures that the first argument to the decorated function is a Tensor.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if isinstance(args[0], Tensor):
+            return func(*args, **kwargs)
+        raise NotImplementedError(f"Function `{func.__name__}` is not implemented for {type(args[0])}")
+
+    return wrapper
+
+
+def dispatch_list(fn: "functools._SingleDispatchCallable", tensor_list: List[Tensor], *args, **kwargs):
+    """
+    Dispatches the function to the type of the wrapped data of the first element in tensor_list.
+
+    :param fn: A function wrapped by `functools.singledispatch`.
+    :param tensor_list: List of Tensors.
+    :return: The result value of the function call.
+    """
+    unwrapped_list = [i.data for i in tensor_list]
+    return fn.dispatch(type(unwrapped_list[0]))(unwrapped_list, *args, **kwargs)
+
+
+def register_numpy_types(singledispatch_fn):
+    """
+    Decorator to register function to singledispatch for numpy classes.
+
+    :param singledispatch_fn: singledispatch function.
+    """
+
+    def inner(func):
+        singledispatch_fn.register(np.ndarray)(func)
+        singledispatch_fn.register(np.generic)(func)
+        return func
+
+    return inner

nncf/experimental/tensor/functions/linalg.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2024 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.
+
+import functools
+from typing import Optional, Tuple, Union
+
+from nncf.experimental.tensor import Tensor
+from nncf.experimental.tensor.functions.dispatcher import tensor_guard
+
+
+@functools.singledispatch
+@tensor_guard
+def norm(
+    a: Tensor,
+    ord: Optional[Union[str, float, int]] = None,
+    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    keepdims: bool = False,
+) -> Tensor:
+    """
+    Computes a vector or matrix norm.
+
+    The following norms can be calculated:
+
+    =====  ============================  ==========================
+    ord    norm for matrices             norm for vectors
+    =====  ============================  ==========================
+    None   Frobenius norm                2-norm
+    'fro'  Frobenius norm                --
+    'nuc'  nuclear norm                  --
+    inf    max(sum(abs(x), axis=1))      max(abs(x))
+    -inf   min(sum(abs(x), axis=1))      min(abs(x))
+    0      --                            sum(x != 0)
+    1      max(sum(abs(x), axis=0))      as below
+    -1     min(sum(abs(x), axis=0))      as below
+    2      2-norm (largest sing. value)  as below
+    -2     smallest singular value       as below
+    other  --                            sum(abs(x)**ord)**(1./ord)
+    =====  ============================  ==========================
+
+    The Frobenius norm is given by [1]_:
+
+        :math:`||A||_F = [\\sum_{i,j} abs(a_{i,j})^2]^{1/2}`
+
+    The nuclear norm is the sum of the singular values.
+
+    Both the Frobenius and nuclear norm orders are only defined for
+    matrices and otherwise raise a ValueError.
+
+    :param a: The input tensor.
+    :param ord: Order of norm. Default: None.
+    :param axis: Axis over which to compute the vector or matrix norm. Default: None.
+    :param keepdims: If set to True, the reduced dimensions are retained in the result
+        as dimensions with size one. Default: False.
+    :return: Norm of the matrix or vector.
+    """
+    return Tensor(norm(a.data, ord, axis, keepdims))