Add Docstring for TF 3x API and Torch 3x Mixed Precision (#1944)

Signed-off-by: zehao-intel <[email protected]>

Author: zehao-intel

.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
+/neural_compressor/torch/algorithms/layer_wise
+/neural_compressor/torch/algorithms/mixed_precision
+/neural_compressor/tensorflow

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:
             <td align="center"><a href="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/tensorflow/tensorflow">TensorFlow</a>/<a href="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/Intel-tensorflow/tensorflow">Intel TensorFlow</a></td>
         </tr>
         <tr>
-            <td rowspan="2" align="center">Smooth Quantization(SQ)</td>
-            <td rowspan="2" align="center">weights</td>
-            <td rowspan="2" align="center">calibration</td>
+            <td rowspan="1" align="center">Smooth Quantization(SQ)</td>
+            <td rowspan="1" align="center">weights</td>
+            <td rowspan="1" align="center">calibration</td>
+            <td align="center">Tensorflow</td>
+            <td align="center"><a href="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/tensorflow/tensorflow">TensorFlow</a>/<a href="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/Intel-tensorflow/tensorflow">Intel TensorFlow</a></td>
+        </tr>
+        <tr>
+            <td rowspan="1" align="center">Mixed Precision(MP)</td>
+            <td rowspan="1" align="center">weights and activations</td>
+            <td rowspan="1" align="center">NA</td>
             <td align="center">Tensorflow</td>
             <td align="center"><a href="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/tensorflow/tensorflow">TensorFlow</a>/<a href="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/Intel-tensorflow/tensorflow">Intel TensorFlow</a></td>
         </tr>
@@ -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.
 

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

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 (

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
 

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

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.
 

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,

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

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 = [