diff --git a/.azure-pipelines/scripts/codeScan/pydocstyle/scan_path.txt b/.azure-pipelines/scripts/codeScan/pydocstyle/scan_path.txt
index c739fe25a47..8eaab8f0b89 100644
--- a/.azure-pipelines/scripts/codeScan/pydocstyle/scan_path.txt
+++ b/.azure-pipelines/scripts/codeScan/pydocstyle/scan_path.txt
@@ -22,4 +22,6 @@
/neural_compressor/torch/export
/neural_compressor/common
/neural_compressor/torch/algorithms/weight_only
-/neural_compressor/torch/algorithms/layer_wise
\ No newline at end of file
+/neural_compressor/torch/algorithms/layer_wise
+/neural_compressor/torch/algorithms/mixed_precision
+/neural_compressor/tensorflow
\ No newline at end of file
diff --git a/docs/source/3x/TensorFlow.md b/docs/source/3x/TensorFlow.md
index dd58c389699..6e4936f2c63 100644
--- a/docs/source/3x/TensorFlow.md
+++ b/docs/source/3x/TensorFlow.md
@@ -2,12 +2,16 @@ TensorFlow
===============
-1. [Introduction](#introduction)
-2. [API for TensorFlow](#api-for-tensorflow)
-3. [Support Matrix](#support-matrix)
- 3.1 [Quantization Scheme](#quantization-scheme)
- 3.2 [Quantization Approaches](#quantization-approaches)
- 3.3 [Backend and Device](#backend-and-device)
+- [TensorFlow](#tensorflow)
+ - [Introduction](#introduction)
+ - [API for TensorFlow](#api-for-tensorflow)
+ - [Support Matrix](#support-matrix)
+ - [Quantization Scheme](#quantization-scheme)
+ - [Quantization Approaches](#quantization-approaches)
+ - [Post Training Static Quantization](#post-training-static-quantization)
+ - [Smooth Quantization](#smooth-quantization)
+ - [Mixed Precision](#mixed-precison)
+ - [Backend and Device](#backend-and-device)
## Introduction
@@ -152,9 +156,16 @@ The supported Quantization methods for TensorFlow and Keras are listed below:
| TensorFlow/Intel TensorFlow |
- | Smooth Quantization(SQ) |
- weights |
- calibration |
+ Smooth Quantization(SQ) |
+ weights |
+ calibration |
+ Tensorflow |
+ TensorFlow/Intel TensorFlow |
+
+
+ | Mixed Precision(MP) |
+ weights and activations |
+ NA |
Tensorflow |
TensorFlow/Intel TensorFlow |
@@ -175,6 +186,10 @@ Smooth Quantization (SQ) is an advanced quantization technique designed to optim
Refer to the [SQ Guide](./TF_SQ.md) for detailed information.
+##### Mixed Precision
+The Mixed Precision (MP) is enabled with Post Training Static Quantization. Once `BF16` is supported on machine, the matched operators will be automatically converted.
+
+
#### Backend and Device
Intel(R) Neural Compressor supports TF GPU with [ITEX-XPU](https://github.com/intel/intel-extension-for-tensorflow). We will automatically run model on GPU by checking if it has been installed.
diff --git a/neural_compressor/tensorflow/algorithms/__init__.py b/neural_compressor/tensorflow/algorithms/__init__.py
index f8d2a4904c2..3373ec76ae7 100644
--- a/neural_compressor/tensorflow/algorithms/__init__.py
+++ b/neural_compressor/tensorflow/algorithms/__init__.py
@@ -11,7 +11,7 @@
# 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.
-
+"""Initialize all algorithms."""
from neural_compressor.tensorflow.algorithms.smoother import SmoothQuant
from neural_compressor.tensorflow.algorithms.static_quant import KerasAdaptor, TensorFlowAdaptor, Tensorflow_ITEXAdaptor
diff --git a/neural_compressor/tensorflow/algorithms/smoother/__init__.py b/neural_compressor/tensorflow/algorithms/smoother/__init__.py
index fd91d612289..60f38f4cbdc 100644
--- a/neural_compressor/tensorflow/algorithms/smoother/__init__.py
+++ b/neural_compressor/tensorflow/algorithms/smoother/__init__.py
@@ -14,6 +14,7 @@
# 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.
+"""Initialize classes of smooth quant algorithm."""
from neural_compressor.tensorflow.algorithms.smoother.core import SmoothQuant
from neural_compressor.tensorflow.algorithms.smoother.scaler import (
diff --git a/neural_compressor/tensorflow/algorithms/smoother/core.py b/neural_compressor/tensorflow/algorithms/smoother/core.py
index 187539ee6eb..64efd596c04 100644
--- a/neural_compressor/tensorflow/algorithms/smoother/core.py
+++ b/neural_compressor/tensorflow/algorithms/smoother/core.py
@@ -14,6 +14,7 @@
# 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.
+"""The core components of sq."""
from typing import Callable, Dict
diff --git a/neural_compressor/tensorflow/algorithms/static_quant/__init__.py b/neural_compressor/tensorflow/algorithms/static_quant/__init__.py
index 05316391218..50d953411fb 100644
--- a/neural_compressor/tensorflow/algorithms/static_quant/__init__.py
+++ b/neural_compressor/tensorflow/algorithms/static_quant/__init__.py
@@ -11,6 +11,7 @@
# 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.
+"""Initialize classes of static quant algorithm."""
from neural_compressor.tensorflow.algorithms.static_quant.keras import KerasAdaptor
from neural_compressor.tensorflow.algorithms.static_quant.tensorflow import TensorFlowAdaptor, Tensorflow_ITEXAdaptor
diff --git a/neural_compressor/tensorflow/algorithms/static_quant/keras.py b/neural_compressor/tensorflow/algorithms/static_quant/keras.py
index 83d9a54609d..c62e92405b8 100644
--- a/neural_compressor/tensorflow/algorithms/static_quant/keras.py
+++ b/neural_compressor/tensorflow/algorithms/static_quant/keras.py
@@ -14,6 +14,7 @@
# 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.
+"""The quantization classes for Keras."""
import copy
import json
@@ -373,9 +374,9 @@ def _calibrate(self, model, dataloader=None, calib_interation=None):
"""Apply calibration.
Args:
- model (tf.keras.Model): The model inserted with FakeQuant layers for calibration.
+ model(tf.keras.Model): The model inserted with FakeQuant layers for calibration.
dataloader(object): The calibration dataloader used to load quantization dataset.
- iteration(int): The iteration of calibration.
+ calib_interation(int): The iteration of calibration.
"""
# run eagerly to fetch the numpy min/max
results = {}
@@ -580,9 +581,9 @@ def _one_shot_query(self):
def _get_specified_version_cfg(self, data):
"""Get the configuration for the current runtime.
+
If there's no matched configuration in the input yaml, we'll
use the `default` field of yaml.
-
Args:
data (Yaml content): input yaml file.
diff --git a/neural_compressor/tensorflow/keras/__init__.py b/neural_compressor/tensorflow/keras/__init__.py
index ab5a1e27d4c..a0539eec3e9 100644
--- a/neural_compressor/tensorflow/keras/__init__.py
+++ b/neural_compressor/tensorflow/keras/__init__.py
@@ -11,6 +11,7 @@
# 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.
+"""Initialize classes for Keras quantization."""
from neural_compressor.tensorflow.keras.quantization import (
StaticQuantConfig,
diff --git a/neural_compressor/tensorflow/keras/layers/__init__.py b/neural_compressor/tensorflow/keras/layers/__init__.py
index 13b83950c97..ee4249138b3 100644
--- a/neural_compressor/tensorflow/keras/layers/__init__.py
+++ b/neural_compressor/tensorflow/keras/layers/__init__.py
@@ -14,6 +14,7 @@
# 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.
+"""Initialize custom layers for Keras quantization."""
from neural_compressor.tensorflow.keras.layers.conv2d import QConv2D
from neural_compressor.tensorflow.keras.layers.dense import QDense
diff --git a/neural_compressor/tensorflow/keras/layers/conv2d.py b/neural_compressor/tensorflow/keras/layers/conv2d.py
index 002ac1c507b..26d6a812ece 100644
--- a/neural_compressor/tensorflow/keras/layers/conv2d.py
+++ b/neural_compressor/tensorflow/keras/layers/conv2d.py
@@ -14,6 +14,7 @@
# 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.
+"""Initialize custom conv2d layers for Keras quantization."""
import json
@@ -33,6 +34,8 @@
if version1_gte_version2(tf.__version__, "2.16.1"): # pragma: no cover
class QConv2D(BaseConv):
+ """The custom quantized Conv2D layer."""
+
def __init__(
self,
name,
@@ -65,6 +68,7 @@ def __init__(
quant_axis=None,
**kwargs
):
+ """Initialize custom quantized Conv2D layer."""
super(QConv2D, self).__init__(
name=name,
rank=2,
@@ -100,6 +104,7 @@ def __init__(
self.quant_axis = quant_axis
def call(self, inputs):
+ """The __call__ function of custom quantized Conv2D layer."""
if self.quant_status == "calib" and not isinstance(inputs, tf.keras.KerasTensor):
if self.granularity == "per_tensor":
self.act_min_value = tf.math.reduce_min(inputs)
@@ -168,9 +173,11 @@ def call(self, inputs):
@classmethod
def from_config(cls, config):
+ """Deserialize this class from a config dict."""
return cls(**config)
def get_config(self):
+ """Serialize this class to a config dict."""
config = super(QConv2D, self).get_config()
config.update(
{
@@ -193,6 +200,8 @@ def get_config(self):
else:
class QConv2D(Conv):
+ """The custom quantized Conv2D layer."""
+
def __init__(
self,
name,
@@ -225,6 +234,7 @@ def __init__(
quant_axis=None,
**kwargs
):
+ """Initialize custom quantized Conv2D layer."""
super(QConv2D, self).__init__(
name=name,
rank=2,
@@ -260,6 +270,7 @@ def __init__(
self.quant_axis = quant_axis
def call(self, inputs):
+ """The __call__ function of custom quantized Conv2D layer."""
if self.quant_status == "calib":
if self.granularity == "per_tensor":
self.act_min_value = tf.math.reduce_min(inputs)
@@ -328,9 +339,11 @@ def call(self, inputs):
@classmethod
def from_config(cls, config):
+ """Deserialize this class from a config dict."""
return cls(**config)
def get_config(self):
+ """Serialize this class to a config dict."""
config = super(QConv2D, self).get_config()
config.update(
{
@@ -352,6 +365,7 @@ def get_config(self):
def initialize_int8_conv2d(fp32_layer, q_config):
+ """Initialize int8 conv2d."""
kwargs = fp32_layer.get_config()
param_list = [
diff --git a/neural_compressor/tensorflow/keras/layers/dense.py b/neural_compressor/tensorflow/keras/layers/dense.py
index 84c4dfabd6c..33bee5a9c78 100644
--- a/neural_compressor/tensorflow/keras/layers/dense.py
+++ b/neural_compressor/tensorflow/keras/layers/dense.py
@@ -14,6 +14,7 @@
# 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.
+"""Initialize custom dense layers for Keras quantization."""
import json
@@ -25,6 +26,8 @@
class QDense(Dense):
+ """The custom quantized Dense layer."""
+
def __init__(
self,
name,
@@ -51,6 +54,7 @@ def __init__(
quant_axis=None,
**kwargs
):
+ """Initialize custom quantized Dense layer."""
super(QDense, self).__init__(
name=name,
units=units,
@@ -79,6 +83,7 @@ def __init__(
self.quant_axis = quant_axis
def call(self, inputs):
+ """The __call__ function of custom quantized Dense layer."""
if self.quant_status == "calib" and not (
version1_gte_version2(tf.__version__, "2.16.1") and isinstance(inputs, tf.keras.KerasTensor)
):
@@ -144,9 +149,11 @@ def call(self, inputs):
@classmethod
def from_config(cls, config):
+ """Deserialize this class from a config dict."""
return cls(**config)
def get_config(self):
+ """Serialize this class to a config dict."""
config = super(QDense, self).get_config()
config.update(
{
@@ -168,6 +175,7 @@ def get_config(self):
def initialize_int8_dense(fp32_layer, q_config):
+ """Initialize int8 dense."""
kwargs = fp32_layer.get_config()
param_list = [
diff --git a/neural_compressor/tensorflow/keras/layers/depthwise_conv2d.py b/neural_compressor/tensorflow/keras/layers/depthwise_conv2d.py
index a0d8511d058..90b823b0fb8 100644
--- a/neural_compressor/tensorflow/keras/layers/depthwise_conv2d.py
+++ b/neural_compressor/tensorflow/keras/layers/depthwise_conv2d.py
@@ -14,6 +14,7 @@
# 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.
+"""Initialize custom depthwise conv2d layers for Keras quantization."""
import json
@@ -35,6 +36,8 @@
if version1_gte_version2(tf.__version__, "2.16.1"):
class QDepthwiseConv2D(BaseDepthwiseConv): # pragma: no cover
+ """The custom quantized DepthwiseConv2D layer."""
+
def __init__(
self,
kernel_size,
@@ -65,6 +68,7 @@ def __init__(
quant_axis=None,
**kwargs
):
+ """Initialize custom quantized DepthwiseConv2D layer."""
super().__init__(
2,
kernel_size=kernel_size,
@@ -98,6 +102,7 @@ def __init__(
self.quant_axis = quant_axis
def call(self, inputs):
+ """The __call__ function of custom quantized DepthwiseConv2D layer."""
if self.quant_status == "calib" and not isinstance(inputs, tf.keras.KerasTensor):
if self.granularity == "per_tensor":
self.act_min_value = tf.math.reduce_min(inputs)
@@ -165,9 +170,11 @@ def call(self, inputs):
@classmethod
def from_config(cls, config):
+ """Deserialize this class from a config dict."""
return cls(**config)
def get_config(self):
+ """Serialize this class to a config dict."""
config = super(QDepthwiseConv2D, self).get_config()
config.update(
{
@@ -190,6 +197,8 @@ def get_config(self):
else:
class QDepthwiseConv2D(DepthwiseConv):
+ """The custom quantized DepthwiseConv2D layer."""
+
def __init__(
self,
kernel_size,
@@ -220,6 +229,7 @@ def __init__(
quant_axis=None,
**kwargs
):
+ """Initialize custom quantized DepthwiseConv2D layer."""
super().__init__(
2,
kernel_size=kernel_size,
@@ -253,6 +263,7 @@ def __init__(
self.quant_axis = quant_axis
def call(self, inputs):
+ """The __call__ function of custom quantized DepthwiseConv2D layer."""
if self.quant_status == "calib":
if self.granularity == "per_tensor":
self.act_min_value = tf.math.reduce_min(inputs)
@@ -315,9 +326,11 @@ def call(self, inputs):
@classmethod
def from_config(cls, config):
+ """Deserialize this class from a config dict."""
return cls(**config)
def get_config(self):
+ """Serialize this class to a config dict."""
config = super(QDepthwiseConv2D, self).get_config()
config.update(
{
@@ -339,6 +352,7 @@ def get_config(self):
@tf_utils.shape_type_conversion
def compute_output_shape(self, input_shape):
+ """Compute_output_shape."""
if self.data_format == "channels_first":
rows = input_shape[2]
cols = input_shape[3]
@@ -369,6 +383,7 @@ def compute_output_shape(self, input_shape):
def initialize_int8_depthwise_conv2d(fp32_layer, q_config):
+ """Initialize int8 depthwise conv2d."""
kwargs = fp32_layer.get_config()
q_name = fp32_layer.name
diff --git a/neural_compressor/tensorflow/keras/layers/layer_initializer.py b/neural_compressor/tensorflow/keras/layers/layer_initializer.py
index d1db0eb3504..c296e24741c 100644
--- a/neural_compressor/tensorflow/keras/layers/layer_initializer.py
+++ b/neural_compressor/tensorflow/keras/layers/layer_initializer.py
@@ -14,6 +14,7 @@
# 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.
+"""Initialize layer initializer functions."""
from neural_compressor.tensorflow.keras.layers.conv2d import initialize_int8_conv2d
from neural_compressor.tensorflow.keras.layers.dense import initialize_int8_dense
diff --git a/neural_compressor/tensorflow/keras/layers/pool2d.py b/neural_compressor/tensorflow/keras/layers/pool2d.py
index 1a04627f06b..5554f8ff709 100644
--- a/neural_compressor/tensorflow/keras/layers/pool2d.py
+++ b/neural_compressor/tensorflow/keras/layers/pool2d.py
@@ -14,6 +14,7 @@
# 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.
+"""Initialize custom pool2d layers for Keras quantization."""
import json
@@ -24,6 +25,8 @@
class QAvgPool2D(AveragePooling2D):
+ """The custom quantized AveragePooling2D layer."""
+
def __init__(
self,
name,
@@ -44,6 +47,7 @@ def __init__(
quant_axis=None,
**kwargs
):
+ """Initialize custom quantized AveragePooling2D layer."""
super(QAvgPool2D, self).__init__(
name=name, pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, **kwargs
)
@@ -61,6 +65,7 @@ def __init__(
self.quant_axis = quant_axis
def __call__(self, inputs):
+ """The __call__ function of custom quantized AveragePooling2D layer."""
if self.quant_status == "calib" and not (
version1_gte_version2(tf.__version__, "2.16.1") and isinstance(inputs, tf.keras.KerasTensor)
):
@@ -94,9 +99,11 @@ def __call__(self, inputs):
@classmethod
def from_config(cls, config):
+ """Deserialize this class from a config dict."""
return cls(**config)
def get_config(self):
+ """Serialize this class to a config dict."""
config = super(QAvgPool2D, self).get_config()
config.update(
{
@@ -118,6 +125,8 @@ def get_config(self):
class QMaxPool2D(MaxPooling2D):
+ """The custom quantized MaxPooling2D layer."""
+
def __init__(
self,
name,
@@ -138,6 +147,7 @@ def __init__(
quant_axis=None,
**kwargs
):
+ """Initialize custom quantized MaxPooling2D layer."""
super(QMaxPool2D, self).__init__(
name=name, pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, **kwargs
)
@@ -155,6 +165,7 @@ def __init__(
self.quant_axis = quant_axis
def __call__(self, inputs):
+ """The __call__ function of custom quantized MaxPooling2D layer."""
if self.quant_status == "calib" and not (
version1_gte_version2(tf.__version__, "2.16.1") and isinstance(inputs, tf.keras.KerasTensor)
):
@@ -189,9 +200,11 @@ def __call__(self, inputs):
@classmethod
def from_config(cls, config):
+ """Deserialize this class from a config dict."""
return cls(**config)
def get_config(self):
+ """Serialize this class to a config dict."""
config = super(QMaxPool2D, self).get_config()
config.update(
{
@@ -213,6 +226,7 @@ def get_config(self):
def initialize_int8_avgpool(fp32_layer, q_config):
+ """Initialize int8 avgpool."""
kwargs = fp32_layer.get_config()
param_list = [
@@ -241,6 +255,7 @@ def initialize_int8_avgpool(fp32_layer, q_config):
def initialize_int8_maxpool(fp32_layer, q_config):
+ """Initialize int8 maxpool."""
kwargs = fp32_layer.get_config()
param_list = [
diff --git a/neural_compressor/tensorflow/keras/layers/separable_conv2d.py b/neural_compressor/tensorflow/keras/layers/separable_conv2d.py
index b3df094fec0..05e23f97fb3 100644
--- a/neural_compressor/tensorflow/keras/layers/separable_conv2d.py
+++ b/neural_compressor/tensorflow/keras/layers/separable_conv2d.py
@@ -14,6 +14,7 @@
# 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.
+"""Initialize custom separable conv2d layers for Keras quantization."""
import json
@@ -35,6 +36,8 @@
if version1_gte_version2(tf.__version__, "2.16.1"): # pragma: no cover
class QSeparableConv2D(BaseSeparableConv):
+ """The custom quantized SeparableConv2D layer."""
+
def __init__(
self,
filters,
@@ -69,6 +72,7 @@ def __init__(
quant_axis=None,
**kwargs
):
+ """Initialize custom quantized SeparableConv2D layer."""
super().__init__(
rank=2,
filters=filters,
@@ -107,6 +111,7 @@ def __init__(
self.quant_axis = quant_axis
def call(self, inputs):
+ """The __call__ function of custom quantized SeparableConv2D layer."""
if self.quant_status == "calib" and not isinstance(inputs, tf.keras.KerasTensor):
if self.granularity == "per_tensor":
self.act_min_value = tf.math.reduce_min(inputs)
@@ -174,9 +179,11 @@ def call(self, inputs):
@classmethod
def from_config(cls, config):
+ """Deserialize this class from a config dict."""
return cls(**config)
def get_config(self):
+ """Serialize this class to a config dict."""
config = super(QSeparableConv2D, self).get_config()
config.update(
{
@@ -199,6 +206,8 @@ def get_config(self):
else:
class QSeparableConv2D(SeparableConv):
+ """The custom quantized SeparableConv2D layer."""
+
def __init__(
self,
filters,
@@ -233,6 +242,7 @@ def __init__(
quant_axis=None,
**kwargs
):
+ """Initialize custom quantized SeparableConv2D layer."""
super().__init__(
rank=2,
filters=filters,
@@ -270,6 +280,7 @@ def __init__(
self.quant_axis = quant_axis
def call(self, inputs):
+ """The __call__ function of custom quantized SeparableConv2D layer."""
if self.quant_status == "calib":
if self.granularity == "per_tensor":
self.act_min_value = tf.math.reduce_min(inputs)
@@ -338,9 +349,11 @@ def call(self, inputs):
@classmethod
def from_config(cls, config):
+ """Deserialize this class from a config dict."""
return cls(**config)
def get_config(self):
+ """Serialize this class to a config dict."""
config = super(QSeparableConv2D, self).get_config()
config.update(
{
@@ -362,6 +375,7 @@ def get_config(self):
def initialize_int8_separable_conv2d(fp32_layer, q_config):
+ """Initialize int8 separable conv2d."""
kwargs = fp32_layer.get_config()
param_list = [
diff --git a/neural_compressor/tensorflow/keras/quantization/__init__.py b/neural_compressor/tensorflow/keras/quantization/__init__.py
index fa0885e3125..4e52426c109 100644
--- a/neural_compressor/tensorflow/keras/quantization/__init__.py
+++ b/neural_compressor/tensorflow/keras/quantization/__init__.py
@@ -11,6 +11,7 @@
# 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.
+"""Initialize configs for Keras quantization."""
from neural_compressor.tensorflow.keras.quantization.config import (
StaticQuantConfig,
diff --git a/neural_compressor/tensorflow/keras/quantization/config.py b/neural_compressor/tensorflow/keras/quantization/config.py
index ae532dc63c4..1cebc23b0e2 100644
--- a/neural_compressor/tensorflow/keras/quantization/config.py
+++ b/neural_compressor/tensorflow/keras/quantization/config.py
@@ -14,6 +14,7 @@
# 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.
+"""The configs of algorithms for Keras."""
from __future__ import annotations
@@ -36,6 +37,8 @@
class OperatorConfig(NamedTuple):
+ """The config for operator."""
+
config: BaseConfig
operators: List[Union[str, Callable]]
valid_func_list: List[Callable] = []
@@ -76,6 +79,7 @@ def __init__(
act_dtype (str): Data type for activations, default is "int8".
act_sym (bool): Indicates whether activations are symmetric, default is True.
act_granularity (str): Calculate tensor-wise scales or channel-wise scales for activations.
+ white_list (list): A list of supported operators of this algorithm.
"""
super().__init__(white_list=white_list)
self.weight_dtype = weight_dtype
@@ -88,6 +92,7 @@ def __init__(
@classmethod
def register_supported_configs(cls) -> List[OperatorConfig]:
+ """Register supported configs."""
supported_configs = []
static_quant_config = StaticQuantConfig(
weight_dtype=["int8", "fp32"],
@@ -112,6 +117,7 @@ def register_supported_configs(cls) -> List[OperatorConfig]:
@staticmethod
def get_model_info(model) -> List[Tuple[str, Callable]]:
+ """Get concrete node names for supported operators."""
white_list = [
"Dense",
"Conv2D",
@@ -133,7 +139,7 @@ def get_model_info(model) -> List[Tuple[str, Callable]]:
@classmethod
def get_config_set_for_tuning(cls) -> Union[None, "StaticQuantConfig", List["StaticQuantConfig"]]:
- # TODO fwk owner needs to update it.
+ """Get a default config set for tuning."""
return StaticQuantConfig(weight_sym=[True, False])
diff --git a/neural_compressor/tensorflow/quantization/algorithm_entry.py b/neural_compressor/tensorflow/quantization/algorithm_entry.py
index e3530bc5e28..5a71e197753 100644
--- a/neural_compressor/tensorflow/quantization/algorithm_entry.py
+++ b/neural_compressor/tensorflow/quantization/algorithm_entry.py
@@ -11,7 +11,7 @@
# 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.
-
+"""The entry interface for algorithms."""
from typing import Callable, Dict
diff --git a/neural_compressor/tensorflow/quantization/autotune.py b/neural_compressor/tensorflow/quantization/autotune.py
index 8dd051b4d38..7c3a1630a54 100644
--- a/neural_compressor/tensorflow/quantization/autotune.py
+++ b/neural_compressor/tensorflow/quantization/autotune.py
@@ -34,6 +34,7 @@
def get_all_config_set() -> Union[BaseConfig, List[BaseConfig]]:
+ """Get all config set."""
return get_all_config_set_from_config_registry(fwk_name=FRAMEWORK_NAME)
diff --git a/neural_compressor/tensorflow/quantization/config.py b/neural_compressor/tensorflow/quantization/config.py
index 738cc61f95a..b4d7fbc701b 100644
--- a/neural_compressor/tensorflow/quantization/config.py
+++ b/neural_compressor/tensorflow/quantization/config.py
@@ -14,7 +14,7 @@
# 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.
-"""Intel Neural Compressor Pytorch quantization config API."""
+"""Intel Neural Compressor TF quantization config API."""
from __future__ import annotations
@@ -93,6 +93,7 @@ def __init__(
act_sym (bool): Indicates whether activations are symmetric, default is True.
act_granularity (str): Calculate tensor-wise scales or channel-wise scales for activations.
act_algorithm (str): Choose quantization algorithms for activations.
+ white_list (list): A list of supported operators of this algorithm.
"""
super().__init__(white_list=white_list)
self.weight_dtype = weight_dtype
@@ -107,6 +108,7 @@ def __init__(
@classmethod
def register_supported_configs(cls) -> List[OperatorConfig]:
+ """Register supported config."""
supported_configs = []
static_quant_config = StaticQuantConfig(
weight_dtype=["int8", "bf16", "fp32"],
@@ -137,6 +139,7 @@ def register_supported_configs(cls) -> List[OperatorConfig]:
cls.supported_configs = supported_configs
def get_model_info(self, model) -> List[Tuple[str, Callable]]:
+ """Get concrete node names for supported operators."""
white_list = [
"Conv2D",
"FusedBatchNormV3",
@@ -168,6 +171,7 @@ def get_model_info(self, model) -> List[Tuple[str, Callable]]:
@classmethod
def get_config_set_for_tuning(cls) -> Union[None, "StaticQuantConfig", List["StaticQuantConfig"]]:
+ """Get a default config set for tuning."""
return StaticQuantConfig(
weight_dtype=["int8", "fp32"],
weight_sym=[True, False],
@@ -230,7 +234,16 @@ def __init__(
"""Init RTN weight-only quantization config.
Args:
- weight_dtype (str): Data type for weights, default is "int".
+ alpha (float): Value to balance input and weight quantization error, between 0 and 1, default is 0.5.
+ folding (bool): Whether to fold mul into the previous layer, default is False.
+ percentile (float): percentile of calibration to remove outliers, default is 99.99.
+ op_types (list): The op types whose input tensor will be dumped, default is ["MatMul", "Conv2D"].
+ scales_per_op (bool): Whether to set individual scale for every op, default is True.
+ record_max_info (bool): whether record the max info in model for alpha tuning, default is False.
+ weight_clip: Whether to clip weight when calculating scales, default is True.
+ auto_alpha_args(dict) : Hyperparameters used to set the alpha search space in SQ auto-tuning,
+ by default the search space is 0.0-1.0 with step_size 0.1.
+ white_list (list): A list of supported operators of this algorithm.
"""
super().__init__()
self.alpha = alpha
@@ -246,6 +259,7 @@ def __init__(
@classmethod
def register_supported_configs(cls) -> List[OperatorConfig]:
+ """Register supported configs."""
supported_configs = []
smooth_quant_config = SmoothQuantConfig()
operators = ["MatMul", "Conv2D"]
@@ -254,6 +268,7 @@ def register_supported_configs(cls) -> List[OperatorConfig]:
@staticmethod
def get_model_info(model) -> List[Tuple[str, Callable]]:
+ """Get concrete node names for supported operators."""
white_list = ["MatMul", "Conv2D"]
filter_result = []
for node in model.graph_def.node:
@@ -265,7 +280,7 @@ def get_model_info(model) -> List[Tuple[str, Callable]]:
@classmethod
def get_config_set_for_tuning(cls) -> Union[None, "SmoothQuantConfig", List["SmoothQuantConfig"]]:
- # TODO fwk owner needs to update it.
+ """Get a default config set for tuning."""
return SmoothQuantConfig(alpha=0.5)
diff --git a/neural_compressor/tensorflow/quantization/quantize.py b/neural_compressor/tensorflow/quantization/quantize.py
index 5a712202dff..95889a6fbe5 100644
--- a/neural_compressor/tensorflow/quantization/quantize.py
+++ b/neural_compressor/tensorflow/quantization/quantize.py
@@ -25,6 +25,7 @@
def need_apply(configs_mapping: Dict[Tuple[str, callable], BaseConfig], algo_name):
+ """Whether to apply the algorithm."""
return any(config.name == algo_name for config in configs_mapping.values())
diff --git a/neural_compressor/tensorflow/quantization/utils/__init__.py b/neural_compressor/tensorflow/quantization/utils/__init__.py
index 28f108cb636..0420d370f04 100644
--- a/neural_compressor/tensorflow/quantization/utils/__init__.py
+++ b/neural_compressor/tensorflow/quantization/utils/__init__.py
@@ -11,3 +11,4 @@
# 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.
+"""The utils for TF quantization."""
diff --git a/neural_compressor/tensorflow/quantization/utils/graph_rewriter/__init__.py b/neural_compressor/tensorflow/quantization/utils/graph_rewriter/__init__.py
index 28f108cb636..c38ca292ee3 100644
--- a/neural_compressor/tensorflow/quantization/utils/graph_rewriter/__init__.py
+++ b/neural_compressor/tensorflow/quantization/utils/graph_rewriter/__init__.py
@@ -11,3 +11,4 @@
# 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.
+"""The utils for rewriting graph."""
diff --git a/neural_compressor/tensorflow/utils/__init__.py b/neural_compressor/tensorflow/utils/__init__.py
index 0e1535b235f..89d5eb8f006 100644
--- a/neural_compressor/tensorflow/utils/__init__.py
+++ b/neural_compressor/tensorflow/utils/__init__.py
@@ -11,6 +11,7 @@
# 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.
+"""The utils for Tensorflow."""
from neural_compressor.tensorflow.utils.data import BaseDataLoader, DummyDataset, DummyDatasetV2
from neural_compressor.tensorflow.utils.model import (
diff --git a/neural_compressor/tensorflow/utils/constants.py b/neural_compressor/tensorflow/utils/constants.py
index c6f2f6aba0c..fc79b116e17 100644
--- a/neural_compressor/tensorflow/utils/constants.py
+++ b/neural_compressor/tensorflow/utils/constants.py
@@ -11,6 +11,7 @@
# 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.
+"""The constants utils for Tensorflow."""
SPR_BASE_VERSIONS = (
"2.11.0202242",
diff --git a/neural_compressor/tensorflow/utils/model.py b/neural_compressor/tensorflow/utils/model.py
index 8ad020678ae..1bf4e6cc259 100644
--- a/neural_compressor/tensorflow/utils/model.py
+++ b/neural_compressor/tensorflow/utils/model.py
@@ -14,6 +14,7 @@
# 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.
+"""The Model API for wrapping all TF model formats."""
import copy
@@ -25,6 +26,8 @@
@singleton
class TensorflowGlobalConfig:
+ """A global config class for setting framework specific information."""
+
global_config = {
"device": "cpu",
"backend": "default",
@@ -36,6 +39,7 @@ class TensorflowGlobalConfig:
}
def reset_global_config(self):
+ """Reset global config with default values."""
self.global_config = copy.deepcopy(TENSORFLOW_DEFAULT_CONFIG)
self.global_config["workspace_path"] = DEFAULT_WORKSPACE
@@ -79,6 +83,7 @@ def __new__(cls, root, **kwargs):
@staticmethod
def set_tf_config(conf, model):
+ """Set tf global config with model information."""
config = TFConfig.global_config
framework = "keras" if isinstance(model, KerasModel) else "tensorflow"
diff --git a/neural_compressor/tensorflow/utils/model_wrappers.py b/neural_compressor/tensorflow/utils/model_wrappers.py
index e1a58f2f53b..6eed7686cc2 100644
--- a/neural_compressor/tensorflow/utils/model_wrappers.py
+++ b/neural_compressor/tensorflow/utils/model_wrappers.py
@@ -14,7 +14,7 @@
# 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.
-"""Class for Tensorflow model."""
+"""The concrete model wrappers for parsing all TF model formats."""
import copy
import datetime
@@ -38,6 +38,7 @@
def get_tf_model_type(model):
+ """The interface of getting type of tensorflow models."""
try:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
@@ -332,6 +333,7 @@ def _get_graph_from_saved_model_v3(model, input_tensor_names, output_tensor_name
model (string or tf.keras.Model): model path or tf.keras.Model object.
input_tensor_names (list of string): input tensor names of the model.
output_tensor_names (list of string): output tensor names of the model.
+
Returns:
graph_def (tf.compat.v1.Session): tf.compat.v1.Session object.
inputs (list of string): validated input names.
@@ -377,6 +379,7 @@ def _get_graph_from_original_keras_v2(model): # pragma: no cover
Args:
model (string or tf.keras.Model): model path or tf.keras.Model object.
+
Returns:
graph_def (tf.compat.v1.Session): tf.compat.v1.Session object.
input_names (list of string): validated input names.
@@ -431,6 +434,7 @@ def _check_keras_format(model, saved_model_dir): # pragma: no cover
Args:
model (string or tf.keras.Model): model path or tf.keras.Model object.
saved_model_dir (string): the path to save a temporary saved_model.
+
Returns:
graph_def (tf.compat.v1.Session): tf.compat.v1.Session object.
inputs (list of string): validated input names.
@@ -460,6 +464,7 @@ def _get_graph_from_saved_model_v1(model):
Args:
model (string or tf.keras.Model): model path or tf.keras.Model object.
+
Returns:
graph_def (tf.compat.v1.Session): tf.compat.v1.Session object.
inputs (list of string): validated input names.
@@ -512,6 +517,7 @@ def try_loading_keras(model, input_tensor_names, output_tensor_names): # pragma
model (string or tf.keras.Model): model path or tf.keras.Model object.
input_tensor_names (list of string): input tensor names of the model.
output_tensor_names (list of string): output tensor names of the model.
+
Returns:
graph_def (tf.compat.v1.Session): tf.compat.v1.Session object.
input_names (list of string): validated input names.
diff --git a/neural_compressor/tensorflow/utils/utility.py b/neural_compressor/tensorflow/utils/utility.py
index 6cbd109bef5..44a7af398a1 100644
--- a/neural_compressor/tensorflow/utils/utility.py
+++ b/neural_compressor/tensorflow/utils/utility.py
@@ -11,6 +11,7 @@
# 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.
+"""The utility functions and classes for Tensorflow."""
import importlib
import logging
@@ -67,6 +68,7 @@ def example_algo(model: tf.keras.Model, quant_config: StaticQuantConfig) -> tf.k
...
Args:
name (str): The name under which the algorithm function will be registered.
+
Returns:
decorator: The decorator function to be used with algorithm functions.
"""
@@ -79,14 +81,16 @@ def decorator(algo_func):
def deep_get(dictionary, keys, default=None):
- """Get the dot key's item in nested dict
- eg person = {'person':{'name':{'first':'John'}}}
- deep_get(person, "person.name.first") will output 'John'.
+ """Get the dot key's item in nested dict.
+ Usage example:
+ person = {'person':{'name':{'first':'John'}}}
+ deep_get(person, "person.name.first") will output 'John'.
Args:
dictionary (dict): The dict object to get keys
keys (dict): The deep keys
default (object): The return item if key not exists
+
Returns:
item: the item of the deep dot keys
"""
diff --git a/neural_compressor/torch/algorithms/mixed_precision/__init__.py b/neural_compressor/torch/algorithms/mixed_precision/__init__.py
index cdd28b3260e..5157e38b3bd 100644
--- a/neural_compressor/torch/algorithms/mixed_precision/__init__.py
+++ b/neural_compressor/torch/algorithms/mixed_precision/__init__.py
@@ -14,6 +14,7 @@
# 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.
+"""Initialize classes for mixed precision."""
from neural_compressor.torch.algorithms.mixed_precision.half_precision_convert import HalfPrecisionConverter
from neural_compressor.torch.algorithms.mixed_precision.module_wrappers import HalfPrecisionModuleWrapper
diff --git a/neural_compressor/torch/quantization/config.py b/neural_compressor/torch/quantization/config.py
index c4b4f6bf04a..29f944b93e3 100644
--- a/neural_compressor/torch/quantization/config.py
+++ b/neural_compressor/torch/quantization/config.py
@@ -1726,6 +1726,8 @@ def __init__(
"""Init MixedPrecision config.
Args:
+ dtype (str or list): The data type of mixed precision, default is fp16.
+ white_list (list): White list of operator names or module types, default is DEFAULT_WHITE_LIST.
"""
super().__init__(white_list=white_list)
self.dtype = dtype
@@ -1733,6 +1735,7 @@ def __init__(
@classmethod
def register_supported_configs(cls) -> List[OperatorConfig]:
+ """Register supported configs."""
supported_configs = []
mixed_precision_config = MixedPrecisionConfig(
dtype=["fp16", "bf16", "fp32"],
@@ -1743,6 +1746,7 @@ def register_supported_configs(cls) -> List[OperatorConfig]:
@staticmethod
def get_model_info(model: torch.nn.Module) -> List[Tuple[str, Callable]]:
+ """Get concrete node names for supported operators."""
white_list = tuple(MixedPrecisionConfig.supported_half_precision_ops)
filter_result = []
for op_name, module in model.named_modules():
@@ -1754,7 +1758,7 @@ def get_model_info(model: torch.nn.Module) -> List[Tuple[str, Callable]]:
@classmethod
def get_config_set_for_tuning(cls) -> Union[None, "MixedPrecisionConfig", List["MixedPrecisionConfig"]]:
- # TODO fwk owner needs to update it.
+ """Get a default config set for tuning."""
return MixedPrecisionConfig(dtype=["fp16", "bf16", "fp32"])