Skip to content

Commit

Permalink
support Pytorch model in the weight compression algorithm
Browse files Browse the repository at this point in the history
  • Loading branch information
alexsu52 committed Jan 9, 2024
1 parent 3621b23 commit 333439d
Show file tree
Hide file tree
Showing 53 changed files with 1,631 additions and 1,019 deletions.
6 changes: 4 additions & 2 deletions docs/api/source/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
6 changes: 4 additions & 2 deletions nncf/common/graph/layer_attributes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
13 changes: 0 additions & 13 deletions nncf/common/graph/transformations/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
18 changes: 9 additions & 9 deletions nncf/experimental/tensor/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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`

Expand All @@ -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>_*",
]
```
6 changes: 3 additions & 3 deletions nncf/experimental/tensor/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -32,6 +33,7 @@ class TensorDataType(Enum):
float64 = auto()
int8 = auto()
uint8 = auto()
int32 = auto()


class TensorDeviceType(Enum):
Expand All @@ -41,3 +43,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
57 changes: 57 additions & 0 deletions nncf/experimental/tensor/functions/__init__.py
Original file line number Diff line number Diff line change
@@ -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.

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 nncf.experimental.tensor.functions.numpy_linalg
import nncf.experimental.tensor.functions.numpy_numeric

try:
import nncf.experimental.tensor.functions.torch_linalg
import nncf.experimental.tensor.functions.torch_numeric # noqa: F401
except ImportError:
pass


_initialize_backends()
57 changes: 57 additions & 0 deletions nncf/experimental/tensor/functions/dispatcher.py
Original file line number Diff line number Diff line change
@@ -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
64 changes: 64 additions & 0 deletions nncf/experimental/tensor/functions/linalg.py
Original file line number Diff line number Diff line change
@@ -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))
Loading

0 comments on commit 333439d

Please sign in to comment.