Skip to content

Commit

Permalink
Inherit OverfloFix and QuantizationPreset from StrEnum (#2608)
Browse files Browse the repository at this point in the history
### Changes

- Inherit `OverflowFix` and `QuantizationPreset` from `StrEnum` instead
of `Enum` to make them serializable.
- Expose `OverflowFix` as `nncf.OverflowFix`

### Reason for changes

Came up during optimum-intel integration.
  • Loading branch information
nikita-savelyevv authored Apr 3, 2024
1 parent 62d1b36 commit 9674f77
Show file tree
Hide file tree
Showing 3 changed files with 6 additions and 3 deletions.
1 change: 1 addition & 0 deletions nncf/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
AdvancedAccuracyRestorerParameters as AdvancedAccuracyRestorerParameters,
)
from nncf.quantization.advanced_parameters import AdvancedQuantizationParameters as AdvancedQuantizationParameters
from nncf.quantization.advanced_parameters import OverflowFix as OverflowFix
from nncf.scopes import IgnoredScope as IgnoredScope
from nncf.scopes import Subgraph as Subgraph
from nncf.version import __version__ as __version__
Expand Down
3 changes: 2 additions & 1 deletion nncf/common/quantization/structs.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from nncf.common.utils.api_marker import api
from nncf.config.schemata.defaults import QUANTIZATION_BITS
from nncf.config.schemata.defaults import QUANTIZATION_PER_CHANNEL
from nncf.parameters import StrEnum
from nncf.parameters import TargetDevice


Expand Down Expand Up @@ -327,7 +328,7 @@ class UnifiedScaleType(Enum):


@api(canonical_alias="nncf.QuantizationPreset")
class QuantizationPreset(Enum):
class QuantizationPreset(StrEnum):
"""
An enum with values corresponding to the available quantization presets.
"""
Expand Down
5 changes: 3 additions & 2 deletions nncf/quantization/advanced_parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,13 +22,14 @@
import nncf
from nncf.common.quantization.structs import QuantizationScheme as QuantizationMode
from nncf.common.utils.api_marker import api
from nncf.parameters import StrEnum
from nncf.quantization.range_estimator import AggregatorType
from nncf.quantization.range_estimator import RangeEstimatorParameters
from nncf.quantization.range_estimator import StatisticsType


@api()
class OverflowFix(Enum):
@api(canonical_alias="nncf.OverflowFix")
class OverflowFix(StrEnum):
"""
This option controls whether to apply the overflow issue fix for the 8-bit
quantization.
Expand Down

0 comments on commit 9674f77

Please sign in to comment.