chore: import upstream snapshot with attribution

2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
@@ -0,0 +1,185 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+#
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#    
+# 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.
+
+"""
+This module is borrowed from TFMOT repository and updated.
+It implements QDQ insertion based on "Last Value Quantization".
+"""
+
+import tensorflow as tf
+
+
+def LastValueQuantize(
+    inputs,
+    min_var,
+    max_var,
+    per_channel=False,
+    channel_axis=-1,
+    name_prefix="LastValueQuant",
+    is_training=True,
+    num_bits=8,
+    narrow_range=False,
+    symmetric=False,
+):
+    """Adds a layer that collects quantization ranges as last input ranges.
+    LastValueQuantize creates variables called 'min' and 'max', representing the
+    interval used for quantization and clamping.
+    Args:
+      inputs: a tensor containing values to be quantized.
+      per_channel: (Optional) a boolean specifying whether to use different
+        quantization ranges per output channel.
+      init_min: a float scalar, the initial value for variable min.
+      init_max: a float scalar, the initial value for variable max.
+      name_prefix: name_prefix for created nodes.
+      is_training: Whether the op is applied to a training or eval graph.
+      num_bits: Number of bits to use for quantization, must be between 2 and 8.
+      narrow_range: Whether to use the narrow quantization range
+        [1; 2^num_bits - 1] or wide range [0; 2^num_bits - 1].
+      symmetric: If true, use symmetric quantization limits instead of training
+        the minimum and maximum of each quantization range separately.
+    Returns:
+      a tensor containing quantized values.
+    """
+    with tf.name_scope(name_prefix):
+        input_shape = inputs.get_shape()
+        input_dim = len(input_shape)
+        if channel_axis == -1:
+            channel_axis += input_dim
+
+        if not is_training:
+            return _QuantizeAndDequantize(
+                inputs,
+                min_var,
+                max_var,
+                per_channel=per_channel,
+                channel_axis=channel_axis,
+                num_bits=num_bits,
+                narrow_range=narrow_range,
+            )
+
+        if per_channel:
+            if input_dim == 2:
+                reduce_dims = [0]
+            elif input_dim == 4:
+                reduce_dims = [i for i in range(input_dim) if i != channel_axis]
+
+        if per_channel:
+            if input_dim >= 2:
+                batch_min = tf.math.reduce_min(
+                    inputs, axis=reduce_dims, name="BatchMin"
+                )
+            else:
+                batch_min = inputs
+        else:
+            batch_min = tf.math.reduce_min(inputs, name="BatchMin")
+
+        if per_channel:
+            if input_dim >= 2:
+                batch_max = tf.math.reduce_max(
+                    inputs, axis=reduce_dims, name="BatchMax"
+                )
+            else:
+                batch_max = inputs
+        else:
+            batch_max = tf.math.reduce_max(inputs, name="BatchMax")
+
+        if symmetric:
+            if narrow_range:
+                min_max_ratio = -1
+            else:
+                # In two's complement notation, the negative range is slightly larger
+                # than the positive range.
+                min_max_ratio = -((1 << num_bits) - 2) / (1 << num_bits)
+
+            # TFLite requires that 0.0 if always in the [min; max] range. Because
+            # batch_min <= batch_max, it follows that range_min <= 0 <= range_max.
+            range_min = tf.math.minimum(batch_min, batch_max / min_max_ratio)
+            range_max = tf.math.maximum(batch_max, batch_min * min_max_ratio)
+        else:
+            # TFLite requires that 0.0 if always in the [min; max] range.
+            range_min = tf.math.minimum(batch_min, 0.0)
+            range_max = tf.math.maximum(batch_max, 0.0)
+
+        assign_min = min_var.assign(range_min, name="AssignMinLast")
+        assign_max = max_var.assign(range_max, name="AssignMaxLast")
+
+        return _QuantizeAndDequantize(
+            inputs,
+            assign_min,
+            assign_max,
+            per_channel=per_channel,
+            channel_axis=channel_axis,
+            num_bits=num_bits,
+            narrow_range=narrow_range,
+        )
+
+
+def _QuantizeAndDequantize(
+    inputs, min_var, max_var, per_channel, channel_axis, num_bits, narrow_range
+):
+    """Adds a fake quantization operation.
+    Depending on value of per_channel, this operation may do global quantization
+    or per channel quantization.  min_var and max_var should have corresponding
+    shapes: [1] when per_channel == False and [d] when per_channel == True.
+    Args:
+      inputs: a tensor containing values to be quantized.
+      min_var: a variable containing quantization range lower end(s).
+      max_var: a variable containing quantization range upper end(s).
+      per_channel: a boolean specifying whether to use per-channel quantization.
+      num_bits: Number of bits to use for quantization, must be between 2 and 8.
+      narrow_range: Whether to use the narrow quantization range
+        [1; 2^num_bits - 1] or wide range [0; 2^num_bits - 1].
+    Returns:
+      a tensor containing quantized values.
+    """
+
+    if per_channel:
+
+        return tf.quantization.quantize_and_dequantize_v2(
+            inputs,
+            min_var,
+            max_var,
+            num_bits=num_bits,
+            narrow_range=narrow_range,
+            axis=channel_axis,
+            range_given=True,
+        )
+    else:
+        assert min_var.get_shape() == []  # pylint: disable=g-explicit-bool-comparison
+        assert max_var.get_shape() == []  # pylint: disable=g-explicit-bool-comparison
+
+        return tf.quantization.quantize_and_dequantize_v2(
+            inputs,
+            min_var,
+            max_var,
+            num_bits=num_bits,
+            narrow_range=narrow_range,
+            range_given=True,
+        )