Merge pull request #209 from fastmachinelearning/feature/opversion_and_trunc_v2

Introduce v1 and v2 for Trunc

Author: maltanar

README.md

@@ -11,7 +11,7 @@
 <img align="left" src="https://xilinx.github.io/finn/img/TFC_1W2A.onnx.png" alt="QONNX example" style="margin-right: 20px" width="200"/>
 
 
-QONNX (Quantized ONNX) introduces several custom operators -- [`IntQuant`](docs/qonnx-custom-ops/intquant_op.md), [`FloatQuant`](docs/qonnx-custom-ops/floatquant_op.md), [`BipolarQuant`](docs/qonnx-custom-ops/bipolar_quant_op.md), and [`Trunc`](docs/qonnx-custom-ops/trunc_op.md) -- in order to represent arbitrary-precision integer and minifloat quantization in ONNX. This enables:
+QONNX (Quantized ONNX) introduces several [custom operators](docs/qonnx-custom-ops/overview.md) -- `IntQuant`, `FloatQuant`, `BipolarQuant`, and `Trunc` -- in order to represent arbitrary-precision integer and minifloat quantization in ONNX. This enables:
 * Representation of binary, ternary, 3-bit, 4-bit, 6-bit or any other integer/fixed-point quantization.
 * Representation of minifloat quantization with configurable exponent and mantissa bits.
 * Quantization is an operator itself, and can be applied to any parameter or layer input.
@@ -29,9 +29,7 @@ This repository contains a set of Python utilities to work with QONNX models, in
 
 ### Operator definitions
 
-* [Quant](docs/qonnx-custom-ops/quant_op.md) for 2-to-arbitrary-bit quantization, with scaling and zero-point
-* [BipolarQuant](docs/qonnx-custom-ops/bipolar_quant_op.md)  for 1-bit (bipolar) quantization, with scaling and zero-point
-* [Trunc](docs/qonnx-custom-ops/trunc_op.md) for truncating to a specified number of bits, with scaling and zero-point
+Please see the [custom operator overview](docs/qonnx-custom-ops/overview.md) table for more details.
 
 ### Installation
 

docs/qonnx-custom-ops/bipolar_quant_op.md renamed to docs/qonnx-custom-ops/bipolarquant_v1.md

@@ -5,7 +5,7 @@ Additionally, takes one float as input, which define the scaling.
 
 #### Version
 
-This operator is not part of the ONNX standard and is not currently versioned.
+The description of this operator in this document corresponds to `qonnx.custom_ops.general` opset version 1.
 
 #### Attributes
 

docs/qonnx-custom-ops/floatquant_op.md renamed to docs/qonnx-custom-ops/floatquant_v1.md

@@ -16,7 +16,7 @@ special (symbolic) values. This makes it nontrivial to infer the maximum represe
 
 #### Version
 
-This operator is not part of the ONNX standard and is not currently versioned.
+The description of this operator in this document corresponds to `qonnx.custom_ops.general` opset version 1.
 
 #### Attributes
 

docs/qonnx-custom-ops/intquant_op.md renamed to docs/qonnx-custom-ops/intquant_v1.md

@@ -9,11 +9,11 @@ rounding_mode defines how quantized values are rounded.
 
 Notes:
 * This operator was previously named `Quant` but is renamed to `IntQuant` to distinguish it from `FloatQuant`. For a transition period, qonnx will transparently handle `Quant` as `IntQuant` for backwards compatibility reasons, but only `IntQuant` should be used for new models.
-* This operator does not work for binary or bipolar quantization, for this purpose the simpler BipolarQuant node exists.
+* This operator does not work for binary or bipolar quantization, for this purpose the simpler `BipolarQuant` node exists.
 
 #### Version
 
-This operator is not part of the ONNX standard and is not currently versioned.
+The description of this operator in this document corresponds to `qonnx.custom_ops.general` opset version 1.
 
 #### Attributes
 

docs/qonnx-custom-ops/overview.md

@@ -0,0 +1,13 @@
+## Operator Schemas
+
+This file lists the QONNX custom operators, similar to `Operators.md` for the ONNX standard.
+It is manually updated, since QONNX custom operators are relatively few in number.
+
+### qonnx.custom_op.general
+
+|**Operator**|**Since version**||
+|-|-|-|
+|<a href="bipolarquant_v1.md">BipolarQuant</a>|<a href="bipolarquant_v1.md">1</a>|
+|<a href="floatquant_v1.md">FloatQuant</a>|<a href="floatquant_v1.md">1</a>|
+|<a href="intquant_v1.md">IntQuant</a>|<a href="intquant_v1.md">1</a>|
+|<a href="trunc_v2.md">Trunc</a>|<a href="trunc_v2.md">2</a>, <a href="trunc_v1.md">1</a>|

docs/qonnx-custom-ops/trunc_op.md renamed to docs/qonnx-custom-ops/trunc_v1.md

@@ -6,7 +6,7 @@ The attribute rounding_mode defines how truncated values are rounded.
 
 #### Version
 
-This operator is not part of the ONNX standard and is not currently versioned.
+The description of this operator in this document corresponds to `qonnx.custom_ops.general` opset version 1.
 
 #### Attributes
 

docs/qonnx-custom-ops/trunc_v2.md

@@ -0,0 +1,144 @@
+### <a name="Trunc"></a><a name="abs">**Trunc**</a>
+
+Truncates the values of one input data (Tensor<T>) at a specified bitwidth and produces one output data (Tensor<T>).
+Additionally, takes four float tensors as input, which define the scale, zero-point, input bit-width and output bit-width of the quantization.
+The attribute rounding_mode defines how truncated values are rounded.
+
+#### Version
+
+This operator is not part of the ONNX standard.
+The description of this operator in this document corresponds to `qonnx.custom_ops.general` opset version 2.
+
+#### Attributes
+
+<dl>
+<dt><tt>rounding_mode</tt> : string (default is "FLOOR")</dt>
+<dd>Defines how rounding should be applied during truncation. Currently available modes are: "ROUND", "CEIL" and "FLOOR". Here "ROUND" implies a round-to-even operation. Lowercase variants for the rounding mode string are also supported: "round", "ceil", "floor".</dd>
+<dt><tt>signed</tt> : int (default is 1)</dt>
+<dd>Defines if the quantization includes a signed bit. E.g. at 8b unsigned=[0, 255] vs signed=[-128, 127].</dd>
+<dt><tt>narrow</tt> : int (default is 0)</dt>
+<dd>Defines if the value range should be interpreted as narrow, when signed=1. E.g. at 8b regular=[-128, 127] vs narrow=[-127, 127].</dd>
+</dl>
+
+#### Inputs
+
+<dl>
+<dt><tt>X</tt> (differentiable) : tensor(float32)</dt>
+<dd>input tensor to truncate</dd>
+<dt><tt>scale</tt> : float32</dt>
+<dd>The scale factor at the input of the truncation</dd>
+<dt><tt>zeropt</tt> : float32</dt>
+<dd>The zero-point at the input of the truncation</dd>
+<dt><tt>in_bitwidth</tt> : int32</dt>
+<dd>The number of bits used at the input of the truncation</dd>
+<dt><tt>out_scale</tt> : float32</dt>
+<dd>The scale factor of the output of the truncation</dd>
+<dt><tt>out_bitwidth</tt> : int32</dt>
+<dd>The number of bits used at the output of the truncation</dd>
+</dl>
+
+
+#### Outputs
+
+<dl>
+<dt><tt>Y</tt> (differentiable) : tensor(float32)</dt>
+<dd>Output tensor</dd>
+</dl>
+
+
+#### Examples
+<details>
+<summary>Trunc</summary>
+
+```python
+from onnx import helper
+import numpy as np
+
+# Define node settings and input
+x = np.random.randn(100).astype(np.float32)*10.
+scale = np.array(1.)
+zeropt = np.array(0.)
+in_bitwidth = np.array(10)
+out_bitwidth = np.array(4)
+rounding_mode = "ROUND"
+
+# Create node
+node = helper.make_node(
+    'Trunc',
+    domain='finn.custom_op.general',
+    inputs=['x', 'scale', 'zeropt', 'in_bitwidth', 'out_bitwidth'],
+    outputs=['y'],
+    rounding_mode=rounding_mode,
+)
+
+# Execute the same settings with the reference implementation (trunc)
+# See the sample implementation for more details on trunc.
+output_ref = trunc(inp_tensor, scale, zeropt, in_bitwidth, out_bitwidth, rounding_mode)
+
+# Execute node and compare
+expect(node, inputs=[x, scale, zeropt, bitwidth], outputs=[output_ref], name='test_trunc')
+
+```
+
+</details>
+
+
+#### Sample Implementation
+
+<details>
+<summary>Trunc</summary>
+
+```python
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import numpy as np
+
+def trunc(inp_tensor, scale, zeropt, input_bit_width, narrow, signed, output_scale, output_bit_width, rounding_mode):
+
+    # Scaling
+    y = inp_tensor / scale
+    y = y + zeropt
+    # Rounding
+    y = np.round(y)
+    # Rescale
+    trunc_scale = 2 ** np.round(
+        np.log2(output_scale / scale)
+    )  # Trunc scale should be a power-of-two - ensure that is the case
+    y = y / trunc_scale
+
+    # Clamping
+    min_int_val = min_int(signed, narrow, output_bit_width)
+    max_int_val = max_int(signed, narrow, output_bit_width)
+    y = np.where(y > max_int_val, max_int_val.astype(y.dtype), y)
+    y = np.where(y < min_int_val, min_int_val.astype(y.dtype), y)
+    # To int (truncate)
+    rounding_fx = resolve_rounding_mode(rounding_mode)
+    y = rounding_fx(y)
+
+    # Rescale
+    output_zeropt = zeropt / trunc_scale  # Rescale zero-point
+    y = y - output_zeropt
+    y = y * output_scale
+
+    return y
+
+def resolve_rounding_mode(mode_string):
+    """Resolve the rounding mode string of Quant and Trunc ops
+    to the corresponding numpy functions."""
+    if mode_string == "ROUND":
+        return np.round
+    elif mode_string == "CEIL":
+        return np.ceil
+    elif mode_string == "FLOOR":
+        return np.floor
+    else:
+        raise ValueError(f"Could not resolve rounding mode called: {mode_string}")
+
+```
+
+</details>

src/qonnx/custom_op/general/__init__.py

@@ -37,7 +37,7 @@
 from qonnx.custom_op.general.multithreshold import MultiThreshold
 from qonnx.custom_op.general.quant import Quant
 from qonnx.custom_op.general.quantavgpool2d import QuantAvgPool2d
-from qonnx.custom_op.general.trunc import Trunc
+from qonnx.custom_op.general.trunc import Trunc_v1, Trunc_v2
 from qonnx.custom_op.general.xnorpopcount import XnorPopcountMatMul
 
 __all__ = [
@@ -51,6 +51,7 @@
     "MultiThreshold",
     "Quant",
     "QuantAvgPool2d",
-    "Trunc",
+    "Trunc_v1",
+    "Trunc_v2",
     "XnorPopcountMatMul",
 ]

src/qonnx/custom_op/general/trunc.py

@@ -31,10 +31,99 @@
 
 from qonnx.core.datatype import DataType
 from qonnx.custom_op.base import CustomOp
-from qonnx.custom_op.general.quant import resolve_rounding_mode
+from qonnx.custom_op.general.quant import max_int, min_int, resolve_rounding_mode
+from qonnx.util.basic import get_preferred_qonnx_opset
 
 
-def trunc(inp_tensor, scale, zeropt, input_bit_width, output_bit_width, rounding_mode):
+def trunc_v2(inp_tensor, scale, zeropt, input_bit_width, narrow, signed, output_scale, output_bit_width, rounding_mode):
+    # Port of TruncIntQuant class from Brevitas: https://bit.ly/3wzIpTR
+
+    # Scaling
+    y = inp_tensor / scale
+    y = y + zeropt
+    # Rounding
+    y = np.round(y)
+    # Rescale
+    trunc_scale = 2 ** np.round(
+        np.log2(output_scale / scale)
+    )  # Trunc scale should be a power-of-two - ensure that is the case
+    y = y / trunc_scale
+
+    # Clamping
+    min_int_val = min_int(signed, narrow, output_bit_width)
+    max_int_val = max_int(signed, narrow, output_bit_width)
+    y = np.where(y > max_int_val, max_int_val.astype(y.dtype), y)
+    y = np.where(y < min_int_val, min_int_val.astype(y.dtype), y)
+    # To int (truncate)
+    rounding_fx = resolve_rounding_mode(rounding_mode)
+    y = rounding_fx(y)
+
+    # Rescale
+    output_zeropt = zeropt / trunc_scale  # Rescale zero-point
+    y = y - output_zeropt
+    y = y * output_scale
+
+    return y
+
+
+class Trunc_v2(CustomOp):
+    """Generic truncation operation for QONNX. Takes four inputs:
+    - input tensor to truncate
+    - the scale
+    - the zero-point
+    - the truncation scale
+    - the truncation bit-width
+
+    The output is a tensor of the same shape as the input tensor, with truncated
+    values.
+    """
+
+    def __init__(self, onnx_node, onnx_opset_version=get_preferred_qonnx_opset()):
+        super().__init__(onnx_node, onnx_opset_version)
+        # override any specified opset version, this instance is v2
+        self.onnx_opset_version = 2
+
+    def get_nodeattr_types(self):
+        return {
+            # The rounding mode, which is used for the trunc function
+            "rounding_mode": ("s", True, "FLOOR"),
+            "narrow": ("i", False, 0, {0, 1}),
+            "signed": ("i", False, 1, {0, 1}),
+        }
+
+    def make_shape_compatible_op(self, model):
+        node = self.onnx_node
+        return helper.make_node("Identity", [node.input[0]], [node.output[0]])
+
+    def infer_node_datatype(self, model):
+        node = self.onnx_node
+        model.set_tensor_datatype(node.output[0], DataType["FLOAT32"])
+
+    def execute_node(self, context, graph):
+        node = self.onnx_node
+        # save inputs
+        inp_tensor = context[node.input[0]]
+        scale = context[node.input[1]]
+        zeropt = context[node.input[2]]
+        input_bit_width = context[node.input[3]]
+        output_scale = context[node.input[4]]
+        output_bit_width = context[node.input[5]]
+        # save attributes
+        rounding_mode = self.get_nodeattr("rounding_mode")
+        narrow = self.get_nodeattr("narrow")
+        signed = self.get_nodeattr("signed")
+        # calculate output
+        ret = trunc_v2(
+            inp_tensor, scale, zeropt, input_bit_width, narrow, signed, output_scale, output_bit_width, rounding_mode
+        )
+        # set context according to output name
+        context[node.output[0]] = ret
+
+    def verify_node(self):
+        pass
+
+
+def trunc_v1(inp_tensor, scale, zeropt, input_bit_width, output_bit_width, rounding_mode):
     # Port of TruncIntQuant class from Brevitas: https://bit.ly/3wzIpTR
 
     # Scaling
@@ -58,7 +147,7 @@ def trunc(inp_tensor, scale, zeropt, input_bit_width, output_bit_width, rounding
     return y
 
 
-class Trunc(CustomOp):
+class Trunc_v1(CustomOp):
     """Generic truncation operation for QONNX. Takes four inputs:
     - input tensor to truncate
     - the scale
@@ -69,6 +158,11 @@ class Trunc(CustomOp):
     values.
     """
 
+    def __init__(self, onnx_node, onnx_opset_version=get_preferred_qonnx_opset()):
+        super().__init__(onnx_node, onnx_opset_version)
+        # override any specified opset version, this instance is v1
+        self.onnx_opset_version = 1
+
     def get_nodeattr_types(self):
         return {
             # The rounding mode, which is used for the trunc function
@@ -94,7 +188,7 @@ def execute_node(self, context, graph):
         # save attributes
         rounding_mode = self.get_nodeattr("rounding_mode")
         # calculate output
-        ret = trunc(inp_tensor, scale, zeropt, input_bit_width, output_bit_width, rounding_mode)
+        ret = trunc_v1(inp_tensor, scale, zeropt, input_bit_width, output_bit_width, rounding_mode)
         # set context according to output name
         context[node.output[0]] = ret