<commit_before># Copyright 2015 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.
# ==============================================================================
r"""Transforms a float-trained graph into an equivalent quantized version.

An example of command-line usage is:
bazel build tensorflow/tools/quantization:quantize_graph \
&& bazel-bin/tensorflow/tools/quantization/quantize_graph \
--input=tensorflow_inception_graph.pb
--output_node_names="softmax2" --print_nodes --output=/tmp/quantized_graph.pb \
--mode=eightbit --logtostderr

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import re
import numpy as np

from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.framework import graph_pb2
from tensorflow.core.framework import node_def_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import graph_util
from tensorflow.python.framework import importer
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import app
from tensorflow.python.platform import flags as flags_lib
from tensorflow.python.platform import gfile

flags = flags_lib
FLAGS = flags.FLAGS

flags.DEFINE_boolean("print_nodes", False, """Lists all nodes in the model.""")
flags.DEFINE_string("input", "", """TensorFlow 'GraphDef' file to load.""")
flags.DEFINE_string("output_node_names", "",
                    """Output node names, comma separated.""")
flags.DEFINE_string("output", "", """File to save the output graph to.""")
flags.DEFINE_integer("bitdepth", 8,
                     """How many bits to quantize the graph to.""")
flags.DEFINE_string("mode", "round",
                    """What transformation to apply (round, quantize,"""
                    """ eightbit, weights, or weights_rounded).""")
flags.DEFINE_string("test_input_dims", "1,224,224,3",
                    """The size of the input tensor to use when testing a"""
                    """ graph loaded from a file.""")
flags.DEFINE_boolean("strip_redundant_quantization", True,
                     """Removes redundant dequantize/quantize pairs.""")
flags.DEFINE_boolean("quantized_input", False,
                     "If true, assume Placeholders are quantized with values "
                     "covering [--quantized_input_min,--quantized_input_max]. "
                     "Only supported when --mode=eightbit")
flags.DEFINE_float("quantized_input_min", 0,
                   "The minimum of the actual input range when "
                   "--quantized_input")
flags.DEFINE_float("quantized_input_max", 1,
                   "The maximum of the actual input range when "
                   "--quantized_input")
flags.DEFINE_float(
    "quantized_fallback_min", None,
    "The fallback 'min' value to use for layers which lack min-max "
    "information. Note: this should be considered a coarse tool just good "
    "enough for experimentation purposes, since graphs quantized in this way "
    "would be very inaccurate.")
flags.DEFINE_float(
    "quantized_fallback_max", None,
    "The fallback 'max' value to use for layers which lack min-max "
    "information. Note: this should be considered a coarse tool just good "
    "enough for experimentation purposes, since graphs quantized in this way "
    "would be very inaccurate.")


def print_input_nodes(current_node, nodes_map, indent, already_visited):
  print(" " * indent + current_node.op + ":" + current_node.name)
  already_visited[current_node.name] = True
  for input_node_name in current_node.input:
    if input_node_name in already_visited:
      continue
    input_node = nodes_map[input_node_name]
    print_input_nodes(input_node, nodes_map, indent + 1, already_visited)


def create_node(op, name, inputs):
  new_node = node_def_pb2.NodeDef()
  new_node.op = op
  new_node.name = name
  for input_name in inputs:
    new_node.input.extend([input_name])
  return new_node


def create_constant_node(name, value, dtype, shape=None):
  node = create_node("Const", name, [])
  set_attr_dtype(node, "dtype", dtype)
  set_attr_tensor(node, "value", value, dtype, shape)
  return node


def copy_attr(node, key, attr_value):
  try:
    node.attr[key].CopyFrom(attr_value)
  except KeyError:
    pass


def set_attr_dtype(node, key, value):
  try:
    node.attr[key].CopyFrom(
        attr_value_pb2.AttrValue(type=value.as_datatype_enum))
  except KeyError:
    pass


def set_attr_shape(node, key, value):
  try:
    node.attr[key].CopyFrom(
        attr_value_pb2.AttrValue(shape=tensor_shape.as_shape(value).as_proto()))
  except KeyError:
    pass


def set_attr_tensor(node, key, value, dtype, shape=None):
  try:
    node.attr[key].CopyFrom(
        attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
            value, dtype=dtype, shape=shape)))
  except KeyError:
    pass


def set_attr_string(node, key, value):
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(s=value))
  except KeyError:
    pass


def set_attr_int_list(node, key, value):
  list_value = attr_value_pb2.AttrValue.ListValue(i=value)
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(list=list_value))
  except KeyError:
    pass


def set_attr_bool(node, key, value):
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(b=value))
  except KeyError:
    pass


def set_attr_int(node, key, value):
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(i=value))
  except KeyError:
    pass


def set_attr_float(node, key, value):
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(f=value))
  except KeyError:
    pass


def node_name_from_input(node_name):
  """Strips off ports and other decorations to get the underlying node name."""
  if node_name.startswith("^"):
    node_name = node_name[1:]
  m = re.search(r"(.*):\d+$", node_name)
  if m:
    node_name = m.group(1)
  return node_name


def ensure_tensor_name_has_port(node_name):
  """Makes sure that a tensor name has :0 if no explicit port exists."""
  m = re.search(r"(.*):\d+$", node_name)
  if m:
    name_with_port = node_name
  else:
    name_with_port = node_name + ":0"
  return name_with_port


def unique_node_name_from_input(node_name):
  """Replaces invalid characters in input names to get a unique node name."""
  return node_name.replace(":", "__port__").replace("^", "__hat__")


def quantize_array(arr, num_buckets):
  """Quantizes a numpy array.

  This function maps each scalar in arr to the center of one of num_buckets
  buckets. For instance,
  quantize_array([0, 0.3, 0.6, 1], 2) => [0.25, 0.25, 0.75, 0.75]

  Args:
    arr: The numpy array to quantize.
    num_buckets: The number of buckets to map "var" to.
  Returns:
    The quantized numpy array.
  Raises:
    ValueError: when num_buckets < 1.
  """
  if num_buckets < 1:
    raise ValueError("num_buckets must be >= 1")
  arr_max = arr.max()
  arr_min = arr.min()
  if arr_max == arr_min:
    return arr
  bucket_width = (arr_max - arr_min) / num_buckets
  # Map scalars to bucket indices. Take special care of max(arr).
  bucket_indices = np.floor((arr - arr_min) / bucket_width)
  bucket_indices[bucket_indices == num_buckets] = num_buckets - 1
  # Map each scalar to the center of a bucket.
  arr = arr_min + bucket_width * (bucket_indices + 0.5)
  return arr


def quantize_weight_rounded(input_node):
  """Returns a replacement node for input_node containing bucketed floats."""
  input_tensor = input_node.attr["value"].tensor
  tensor_value = tensor_util.MakeNdarray(input_tensor)
  shape = input_tensor.tensor_shape
  # Currently, the parameter FLAGS.bitdepth is used to compute the
  # number of buckets as 1 << FLAGS.bitdepth, meaning the number of
  # buckets can only be a power of 2.
  # This could be fixed by introducing a new parameter, num_buckets,
  # which would allow for more flexibility in chosing the right model
  # size/accuracy tradeoff. But I didn't want to add more parameters
  # to this script than absolutely necessary.
  num_buckets = 1 << FLAGS.bitdepth
  tensor_value_rounded = quantize_array(tensor_value, num_buckets)
  tensor_shape_list = tensor_util.TensorShapeProtoToList(shape)
  return [
      create_constant_node(
          input_node.name,
          tensor_value_rounded,
          dtypes.float32,
          shape=tensor_shape_list)
  ]


def quantize_weight_eightbit(input_node, quantization_mode):
  """Returns replacement nodes for input_node using the Dequantize op."""
  base_name = input_node.name + "_"
  quint8_const_name = base_name + "quint8_const"
  min_name = base_name + "min"
  max_name = base_name + "max"
  float_tensor = tensor_util.MakeNdarray(input_node.attr["value"].tensor)
  min_value = np.min(float_tensor.flatten())
  max_value = np.max(float_tensor.flatten())
  # Make sure that the range includes zero.
  if min_value > 0.0:
    min_value = 0.0
  # min_value == max_value is a tricky case. It can occur for general
  # tensors, and of course for scalars. The quantized ops cannot deal
  # with this case, so we set max_value to something else.
  # It's a tricky question what is the numerically best solution to
  # deal with this degeneracy.
  # TODO(petewarden): Better use a tolerance than a hard comparison?
  if min_value == max_value:
    if abs(min_value) < 0.000001:
      max_value = min_value + 1.0
    elif min_value > 0:
      max_value = 2 * min_value
    else:
      max_value = min_value / 2.0

  sess = session.Session()
  with sess.as_default():
    quantize_op = array_ops.quantize_v2(
        float_tensor,
        min_value,
        max_value,
        dtypes.quint8,
        mode=quantization_mode)
    quint8_tensor = quantize_op[0].eval()
  shape = tensor_util.TensorShapeProtoToList(input_node.attr["value"]
                                             .tensor.tensor_shape)
  quint8_const_node = create_constant_node(
      quint8_const_name, quint8_tensor, dtypes.quint8, shape=shape)
  min_node = create_constant_node(min_name, min_value, dtypes.float32)
  max_node = create_constant_node(max_name, max_value, dtypes.float32)
  dequantize_node = create_node("Dequantize", input_node.name,
                                [quint8_const_name, min_name, max_name])
  set_attr_dtype(dequantize_node, "T", dtypes.quint8)
  set_attr_string(dequantize_node, "mode", quantization_mode)
  return [quint8_const_node, min_node, max_node, dequantize_node]


EightbitizeRecursionState = collections.namedtuple(
    "EightbitizeRecursionState",
    ["already_visited", "output_node_stack", "merged_with_fake_quant"])


class GraphRewriter(object):
  """Takes a float graph, and rewrites it in quantized form."""

  def __init__(self,
               input_graph,
               mode,
               quantized_input_range,
               fallback_quantization_range=None):
    """Sets up the class to rewrite a float graph.

    Args:
      input_graph: A float graph to transform.
      mode: A string controlling how quantization is performed -
        round, quantize, eightbit, or weights.
      quantized_input_range: if set, assume the input is
        quantized and represents the range
        [quantized_input_range[0], quantized_input_range[1]]
      fallback_quantization_range: if set, then for nodes where the quantization
        range can't be inferred from the graph, use the range
        [fallback_quantization_range[0], fallback_quantization_range[1]) instead
        of using a RequantizationRange node in the graph.

    Raises:
      ValueError: Two nodes with the same name were found in the graph.
    """
    self.input_graph = input_graph
    self.nodes_map = self.create_nodes_map(input_graph)
    self.output_graph = None
    self.mode = mode
    self.final_node_renames = {}
    if quantized_input_range:
      self.input_range = (quantized_input_range[0], quantized_input_range[1])
      if self.input_range[0] >= self.input_range[1]:
        raise ValueError("Invalid quantized_input_range: [%s,%s]" %
                         self.input_range)
      if self.mode != "eightbit":
        raise ValueError(
            "quantized_input_range can only be specified in eightbit mode")
    else:
      self.input_range = None

    if fallback_quantization_range:
      self.fallback_quantization_range = [
          fallback_quantization_range[0], fallback_quantization_range[1]
      ]
      if (self.fallback_quantization_range[0] >=
          self.fallback_quantization_range[1]):
        raise ValueError("Invalid fallback_quantization_range: [%s,%s]" %
                         self.fallback_quantization_range)
      if self.mode != "eightbit":
        raise ValueError("fallback_quantization_range can only be "
                         "specified in eightbit mode")
    else:
      self.fallback_quantization_range = None

    # Data that is valid only during the recursive call to rewrite the graph.
    self.state = None

  def create_nodes_map(self, graph):
    """Builds a mapping of node names to their defs from the graph."""
    nodes_map = {}
    for node in graph.node:
      if node.name not in nodes_map.keys():
        nodes_map[node.name] = node
      else:
        raise ValueError("Duplicate node names detected.")
    return nodes_map

  def rewrite(self, output_node_names):
    """Triggers rewriting of the float graph.

    Args:
      output_node_names: A list of names of the nodes that produce the final
        results.

    Returns:
      A quantized version of the float graph.
    """
    self.output_graph = graph_pb2.GraphDef()
    output_nodes = [
        self.nodes_map[output_node_name]
        for output_node_name in output_node_names
    ]
    if self.mode == "round":
      self.already_visited = {}
      for output_node in output_nodes:
        self.round_nodes_recursively(output_node)
    elif self.mode == "quantize":
      self.already_visited = {}
      self.already_quantized = {}
      for output_node in output_nodes:
        self.quantize_nodes_recursively(output_node)
    elif self.mode == "eightbit":
      self.set_input_graph(graph_util.remove_training_nodes(self.input_graph))
      output_nodes = [
          self.nodes_map[output_node_name]
          for output_node_name in output_node_names
      ]

      self.state = EightbitizeRecursionState(
          already_visited={}, output_node_stack=[], merged_with_fake_quant={})
      for output_node in output_nodes:
        self.eightbitize_nodes_recursively(output_node)
      self.state = None
      if self.input_range:
        self.add_output_graph_node(
            create_constant_node("quantized_input_min_value", self.input_range[
                0], dtypes.float32, []))
        self.add_output_graph_node(
            create_constant_node("quantized_input_max_value", self.input_range[
                1], dtypes.float32, []))
      if self.fallback_quantization_range:
        self.add_output_graph_node(
            create_constant_node("fallback_quantization_min_value",
                                 self.fallback_quantization_range[0],
                                 dtypes.float32, []))
        self.add_output_graph_node(
            create_constant_node("fallback_quantization_max_value",
                                 self.fallback_quantization_range[1],
                                 dtypes.float32, []))
      if FLAGS.strip_redundant_quantization:
        self.output_graph = self.remove_redundant_quantization(
            self.output_graph)
        self.remove_dead_nodes(output_node_names)
      self.apply_final_node_renames()
    elif self.mode == "weights":
      self.output_graph = self.quantize_weights(self.input_graph,
                                                b"MIN_COMBINED")
      self.remove_dead_nodes(output_node_names)
    elif self.mode == "weights_rounded":
      self.output_graph = self.quantize_weights(self.input_graph, self.mode)
      self.remove_dead_nodes(output_node_names)
    else:
      print("Bad mode - " + self.mode + ".")
    return self.output_graph

  def round_nodes_recursively(self, current_node):
    """The entry point for simple rounding quantization."""
    if self.already_visited[current_node.name]:
      return
    self.already_visited[current_node.name] = True
    for input_node_name in current_node.input:
      input_node_name = node_name_from_input(input_node_name)
      input_node = self.nodes_map[input_node_name]
      self.round_nodes_recursively(input_node)
    nodes_to_quantize = ["Conv2D", "BiasAdd", "MatMul"]
    if any(current_node.op in s for s in nodes_to_quantize):
      new_node = node_def_pb2.NodeDef()
      new_node.CopyFrom(current_node)
      new_node.name = current_node.name + "_original"
      self.add_output_graph_node(new_node)
      levels = 1 << FLAGS.bitdepth
      constant_name = current_node.name + "_round_depth"
      constant_tensor = constant_op.constant(
          levels, dtype=dtypes.int32, name=constant_name)
      constant_node = constant_tensor.op.node_def<commit_msg>Quantize Graph - check whether the key exists before using it (#9393)<commit_after># Copyright 2015 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.
# ==============================================================================
r"""Transforms a float-trained graph into an equivalent quantized version.

An example of command-line usage is:
bazel build tensorflow/tools/quantization:quantize_graph \
&& bazel-bin/tensorflow/tools/quantization/quantize_graph \
--input=tensorflow_inception_graph.pb
--output_node_names="softmax2" --print_nodes --output=/tmp/quantized_graph.pb \
--mode=eightbit --logtostderr

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import re
import numpy as np

from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.framework import graph_pb2
from tensorflow.core.framework import node_def_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import graph_util
from tensorflow.python.framework import importer
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import app
from tensorflow.python.platform import flags as flags_lib
from tensorflow.python.platform import gfile

flags = flags_lib
FLAGS = flags.FLAGS

flags.DEFINE_boolean("print_nodes", False, """Lists all nodes in the model.""")
flags.DEFINE_string("input", "", """TensorFlow 'GraphDef' file to load.""")
flags.DEFINE_string("output_node_names", "",
                    """Output node names, comma separated.""")
flags.DEFINE_string("output", "", """File to save the output graph to.""")
flags.DEFINE_integer("bitdepth", 8,
                     """How many bits to quantize the graph to.""")
flags.DEFINE_string("mode", "round",
                    """What transformation to apply (round, quantize,"""
                    """ eightbit, weights, or weights_rounded).""")
flags.DEFINE_string("test_input_dims", "1,224,224,3",
                    """The size of the input tensor to use when testing a"""
                    """ graph loaded from a file.""")
flags.DEFINE_boolean("strip_redundant_quantization", True,
                     """Removes redundant dequantize/quantize pairs.""")
flags.DEFINE_boolean("quantized_input", False,
                     "If true, assume Placeholders are quantized with values "
                     "covering [--quantized_input_min,--quantized_input_max]. "
                     "Only supported when --mode=eightbit")
flags.DEFINE_float("quantized_input_min", 0,
                   "The minimum of the actual input range when "
                   "--quantized_input")
flags.DEFINE_float("quantized_input_max", 1,
                   "The maximum of the actual input range when "
                   "--quantized_input")
flags.DEFINE_float(
    "quantized_fallback_min", None,
    "The fallback 'min' value to use for layers which lack min-max "
    "information. Note: this should be considered a coarse tool just good "
    "enough for experimentation purposes, since graphs quantized in this way "
    "would be very inaccurate.")
flags.DEFINE_float(
    "quantized_fallback_max", None,
    "The fallback 'max' value to use for layers which lack min-max "
    "information. Note: this should be considered a coarse tool just good "
    "enough for experimentation purposes, since graphs quantized in this way "
    "would be very inaccurate.")


def print_input_nodes(current_node, nodes_map, indent, already_visited):
  print(" " * indent + current_node.op + ":" + current_node.name)
  already_visited[current_node.name] = True
  for input_node_name in current_node.input:
    if input_node_name in already_visited:
      continue
    input_node = nodes_map[input_node_name]
    print_input_nodes(input_node, nodes_map, indent + 1, already_visited)


def create_node(op, name, inputs):
  new_node = node_def_pb2.NodeDef()
  new_node.op = op
  new_node.name = name
  for input_name in inputs:
    new_node.input.extend([input_name])
  return new_node


def create_constant_node(name, value, dtype, shape=None):
  node = create_node("Const", name, [])
  set_attr_dtype(node, "dtype", dtype)
  set_attr_tensor(node, "value", value, dtype, shape)
  return node


def copy_attr(node, key, attr_value):
  try:
    node.attr[key].CopyFrom(attr_value)
  except KeyError:
    pass


def set_attr_dtype(node, key, value):
  try:
    node.attr[key].CopyFrom(
        attr_value_pb2.AttrValue(type=value.as_datatype_enum))
  except KeyError:
    pass


def set_attr_shape(node, key, value):
  try:
    node.attr[key].CopyFrom(
        attr_value_pb2.AttrValue(shape=tensor_shape.as_shape(value).as_proto()))
  except KeyError:
    pass


def set_attr_tensor(node, key, value, dtype, shape=None):
  try:
    node.attr[key].CopyFrom(
        attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
            value, dtype=dtype, shape=shape)))
  except KeyError:
    pass


def set_attr_string(node, key, value):
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(s=value))
  except KeyError:
    pass


def set_attr_int_list(node, key, value):
  list_value = attr_value_pb2.AttrValue.ListValue(i=value)
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(list=list_value))
  except KeyError:
    pass


def set_attr_bool(node, key, value):
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(b=value))
  except KeyError:
    pass


def set_attr_int(node, key, value):
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(i=value))
  except KeyError:
    pass


def set_attr_float(node, key, value):
  try:
    node.attr[key].CopyFrom(attr_value_pb2.AttrValue(f=value))
  except KeyError:
    pass


def node_name_from_input(node_name):
  """Strips off ports and other decorations to get the underlying node name."""
  if node_name.startswith("^"):
    node_name = node_name[1:]
  m = re.search(r"(.*):\d+$", node_name)
  if m:
    node_name = m.group(1)
  return node_name


def ensure_tensor_name_has_port(node_name):
  """Makes sure that a tensor name has :0 if no explicit port exists."""
  m = re.search(r"(.*):\d+$", node_name)
  if m:
    name_with_port = node_name
  else:
    name_with_port = node_name + ":0"
  return name_with_port


def unique_node_name_from_input(node_name):
  """Replaces invalid characters in input names to get a unique node name."""
  return node_name.replace(":", "__port__").replace("^", "__hat__")


def quantize_array(arr, num_buckets):
  """Quantizes a numpy array.

  This function maps each scalar in arr to the center of one of num_buckets
  buckets. For instance,
  quantize_array([0, 0.3, 0.6, 1], 2) => [0.25, 0.25, 0.75, 0.75]

  Args:
    arr: The numpy array to quantize.
    num_buckets: The number of buckets to map "var" to.
  Returns:
    The quantized numpy array.
  Raises:
    ValueError: when num_buckets < 1.
  """
  if num_buckets < 1:
    raise ValueError("num_buckets must be >= 1")
  arr_max = arr.max()
  arr_min = arr.min()
  if arr_max == arr_min:
    return arr
  bucket_width = (arr_max - arr_min) / num_buckets
  # Map scalars to bucket indices. Take special care of max(arr).
  bucket_indices = np.floor((arr - arr_min) / bucket_width)
  bucket_indices[bucket_indices == num_buckets] = num_buckets - 1
  # Map each scalar to the center of a bucket.
  arr = arr_min + bucket_width * (bucket_indices + 0.5)
  return arr


def quantize_weight_rounded(input_node):
  """Returns a replacement node for input_node containing bucketed floats."""
  input_tensor = input_node.attr["value"].tensor
  tensor_value = tensor_util.MakeNdarray(input_tensor)
  shape = input_tensor.tensor_shape
  # Currently, the parameter FLAGS.bitdepth is used to compute the
  # number of buckets as 1 << FLAGS.bitdepth, meaning the number of
  # buckets can only be a power of 2.
  # This could be fixed by introducing a new parameter, num_buckets,
  # which would allow for more flexibility in chosing the right model
  # size/accuracy tradeoff. But I didn't want to add more parameters
  # to this script than absolutely necessary.
  num_buckets = 1 << FLAGS.bitdepth
  tensor_value_rounded = quantize_array(tensor_value, num_buckets)
  tensor_shape_list = tensor_util.TensorShapeProtoToList(shape)
  return [
      create_constant_node(
          input_node.name,
          tensor_value_rounded,
          dtypes.float32,
          shape=tensor_shape_list)
  ]


def quantize_weight_eightbit(input_node, quantization_mode):
  """Returns replacement nodes for input_node using the Dequantize op."""
  base_name = input_node.name + "_"
  quint8_const_name = base_name + "quint8_const"
  min_name = base_name + "min"
  max_name = base_name + "max"
  float_tensor = tensor_util.MakeNdarray(input_node.attr["value"].tensor)
  min_value = np.min(float_tensor.flatten())
  max_value = np.max(float_tensor.flatten())
  # Make sure that the range includes zero.
  if min_value > 0.0:
    min_value = 0.0
  # min_value == max_value is a tricky case. It can occur for general
  # tensors, and of course for scalars. The quantized ops cannot deal
  # with this case, so we set max_value to something else.
  # It's a tricky question what is the numerically best solution to
  # deal with this degeneracy.
  # TODO(petewarden): Better use a tolerance than a hard comparison?
  if min_value == max_value:
    if abs(min_value) < 0.000001:
      max_value = min_value + 1.0
    elif min_value > 0:
      max_value = 2 * min_value
    else:
      max_value = min_value / 2.0

  sess = session.Session()
  with sess.as_default():
    quantize_op = array_ops.quantize_v2(
        float_tensor,
        min_value,
        max_value,
        dtypes.quint8,
        mode=quantization_mode)
    quint8_tensor = quantize_op[0].eval()
  shape = tensor_util.TensorShapeProtoToList(input_node.attr["value"]
                                             .tensor.tensor_shape)
  quint8_const_node = create_constant_node(
      quint8_const_name, quint8_tensor, dtypes.quint8, shape=shape)
  min_node = create_constant_node(min_name, min_value, dtypes.float32)
  max_node = create_constant_node(max_name, max_value, dtypes.float32)
  dequantize_node = create_node("Dequantize", input_node.name,
                                [quint8_const_name, min_name, max_name])
  set_attr_dtype(dequantize_node, "T", dtypes.quint8)
  set_attr_string(dequantize_node, "mode", quantization_mode)
  return [quint8_const_node, min_node, max_node, dequantize_node]


EightbitizeRecursionState = collections.namedtuple(
    "EightbitizeRecursionState",
    ["already_visited", "output_node_stack", "merged_with_fake_quant"])


class GraphRewriter(object):
  """Takes a float graph, and rewrites it in quantized form."""

  def __init__(self,
               input_graph,
               mode,
               quantized_input_range,
               fallback_quantization_range=None):
    """Sets up the class to rewrite a float graph.

    Args:
      input_graph: A float graph to transform.
      mode: A string controlling how quantization is performed -
        round, quantize, eightbit, or weights.
      quantized_input_range: if set, assume the input is
        quantized and represents the range
        [quantized_input_range[0], quantized_input_range[1]]
      fallback_quantization_range: if set, then for nodes where the quantization
        range can't be inferred from the graph, use the range
        [fallback_quantization_range[0], fallback_quantization_range[1]) instead
        of using a RequantizationRange node in the graph.

    Raises:
      ValueError: Two nodes with the same name were found in the graph.
    """
    self.input_graph = input_graph
    self.nodes_map = self.create_nodes_map(input_graph)
    self.output_graph = None
    self.mode = mode
    self.final_node_renames = {}
    if quantized_input_range:
      self.input_range = (quantized_input_range[0], quantized_input_range[1])
      if self.input_range[0] >= self.input_range[1]:
        raise ValueError("Invalid quantized_input_range: [%s,%s]" %
                         self.input_range)
      if self.mode != "eightbit":
        raise ValueError(
            "quantized_input_range can only be specified in eightbit mode")
    else:
      self.input_range = None

    if fallback_quantization_range:
      self.fallback_quantization_range = [
          fallback_quantization_range[0], fallback_quantization_range[1]
      ]
      if (self.fallback_quantization_range[0] >=
          self.fallback_quantization_range[1]):
        raise ValueError("Invalid fallback_quantization_range: [%s,%s]" %
                         self.fallback_quantization_range)
      if self.mode != "eightbit":
        raise ValueError("fallback_quantization_range can only be "
                         "specified in eightbit mode")
    else:
      self.fallback_quantization_range = None

    # Data that is valid only during the recursive call to rewrite the graph.
    self.state = None

  def create_nodes_map(self, graph):
    """Builds a mapping of node names to their defs from the graph."""
    nodes_map = {}
    for node in graph.node:
      if node.name not in nodes_map.keys():
        nodes_map[node.name] = node
      else:
        raise ValueError("Duplicate node names detected.")
    return nodes_map

  def rewrite(self, output_node_names):
    """Triggers rewriting of the float graph.

    Args:
      output_node_names: A list of names of the nodes that produce the final
        results.

    Returns:
      A quantized version of the float graph.
    """
    self.output_graph = graph_pb2.GraphDef()
    output_nodes = [
        self.nodes_map[output_node_name]
        for output_node_name in output_node_names
    ]
    if self.mode == "round":
      self.already_visited = {}
      for output_node in output_nodes:
        self.round_nodes_recursively(output_node)
    elif self.mode == "quantize":
      self.already_visited = {}
      self.already_quantized = {}
      for output_node in output_nodes:
        self.quantize_nodes_recursively(output_node)
    elif self.mode == "eightbit":
      self.set_input_graph(graph_util.remove_training_nodes(self.input_graph))
      output_nodes = [
          self.nodes_map[output_node_name]
          for output_node_name in output_node_names
      ]

      self.state = EightbitizeRecursionState(
          already_visited={}, output_node_stack=[], merged_with_fake_quant={})
      for output_node in output_nodes:
        self.eightbitize_nodes_recursively(output_node)
      self.state = None
      if self.input_range:
        self.add_output_graph_node(
            create_constant_node("quantized_input_min_value", self.input_range[
                0], dtypes.float32, []))
        self.add_output_graph_node(
            create_constant_node("quantized_input_max_value", self.input_range[
                1], dtypes.float32, []))
      if self.fallback_quantization_range:
        self.add_output_graph_node(
            create_constant_node("fallback_quantization_min_value",
                                 self.fallback_quantization_range[0],
                                 dtypes.float32, []))
        self.add_output_graph_node(
            create_constant_node("fallback_quantization_max_value",
                                 self.fallback_quantization_range[1],
                                 dtypes.float32, []))
      if FLAGS.strip_redundant_quantization:
        self.output_graph = self.remove_redundant_quantization(
            self.output_graph)
        self.remove_dead_nodes(output_node_names)
      self.apply_final_node_renames()
    elif self.mode == "weights":
      self.output_graph = self.quantize_weights(self.input_graph,
                                                b"MIN_COMBINED")
      self.remove_dead_nodes(output_node_names)
    elif self.mode == "weights_rounded":
      self.output_graph = self.quantize_weights(self.input_graph, self.mode)
      self.remove_dead_nodes(output_node_names)
    else:
      print("Bad mode - " + self.mode + ".")
    return self.output_graph

  def round_nodes_recursively(self, current_node):
    """The entry point for simple rounding quantization."""
    if (current_node.name in self.already_visited) and self.already_visited[current_node.name]:
      return
    self.already_visited[current_node.name] = True
    for input_node_name in current_node.input:
      input_node_name = node_name_from_input(input_node_name)
      input_node = self.nodes_map[input_node_name]
      self.round_nodes_recursively(input_node)
    nodes_to_quantize = ["Conv2D", "BiasAdd", "MatMul"]
    if any(current_node.op in s for s in nodes_to_quantize):
      new_node = node_def_pb2.NodeDef()
      new_node.CopyFrom(current_node)
      new_node.name = current_node.name + "_original"
      self.add_output_graph_node(new_node)
      levels = 1 << FLAGS.bitdepth
      constant_name = current_node.name + "_round_depth"
      constant_tensor = constant_op.constant(
          levels, dtype=dtypes.int32, name=constant_name)
      constant_node = constant_tensor.op.node_def
{'fraction_non_alphanumeric': 0.09435972191074232, 'fraction_numerical': 0.007512895267997309, 'mean_word_length': 3.3636697247706424, 'pattern_counts': {'":': 20, '<': 21, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6979396', 'n_tokens_mistral': 11171, 'n_tokens_neox': 10747, 'n_words': 2809}