Spaces:

Kasamuday
/

object

Runtime error

object / models /research /seq_flow_lite /layers /quantization_layers.py

Upload 4260 files

f3507ef verified over 2 years ago

4.34 kB

# Copyright 2020 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.
# ==============================================================================
"""Layers for quantization."""

import tensorflow as tf

from layers import base_layers # import seq_flow_lite module


class ActivationQuantization(base_layers.BaseLayer):
  """A class that applies quantization to a activation tensor."""

  def __init__(self, ema_decay=0.99, num_bits=8, **kwargs):
    self.ema_decay = ema_decay
    self.num_bits = num_bits
    super(ActivationQuantization, self).__init__(**kwargs)

  def build(self, input_shapes):
    if self.parameters.quantize:
      self.min_var = self.add_weight(
          "min", initializer=tf.keras.initializers.Zeros(), trainable=False)
      self.max_var = self.add_weight(
          "max", initializer=tf.keras.initializers.Ones(), trainable=False)

  def call(self, inputs):
    if self.parameters.quantize:
      if self.parameters.mode == base_layers.TRAIN:
        # Toco expects 0.0 to be part of the quantization range.
        batch_min = tf.minimum(tf.reduce_min(inputs), 0.0)
        min_var = self.assign_moving_average(self.min_var, batch_min,
                                             self.ema_decay)

        batch_max = tf.maximum(tf.reduce_max(inputs), 0.0)
        max_var = self.assign_moving_average(self.max_var, batch_max,
                                             self.ema_decay)
        with tf.control_dependencies([min_var, max_var]):
          return tf.quantization.fake_quant_with_min_max_vars(
              inputs, batch_min, batch_max, num_bits=self.num_bits)
      else:
        return tf.quantization.fake_quant_with_min_max_vars(
            inputs, self.min_var, self.max_var, num_bits=self.num_bits)
    return inputs

  def quantize_using_range(self, inputs):
    # This method can only be called after a call to "call" method in this class
    if self.parameters.quantize:
      return tf.quantization.fake_quant_with_min_max_vars(
          inputs, self.min_var, self.max_var, num_bits=self.num_bits)
    return inputs


class ConcatQuantization(ActivationQuantization):
  """A class that applies quantization to a activation tensor."""

  def __init__(self, axis=2, **kwargs):
    self.axis = axis
    super(ConcatQuantization, self).__init__(**kwargs)

  def _reduce_list(self, tensor_list, functor):
    reduce_result = [functor(tensor) for tensor in tensor_list]
    # Toco expects 0.0 to be part of the quantization range.
    reduce_result.append(tf.constant(0.0))
    return functor(tf.stack(reduce_result))

  def call(self, tensors):
    # Ignore empty invocations done to build the keras layer.
    if tensors is None:
      return
    if self.parameters.quantize:
      if self.parameters.mode == base_layers.TRAIN:
        # Toco expects 0.0 to be part of the quantization range.
        batch_min = self._reduce_list(tensors, tf.reduce_min)
        min_var = self.assign_moving_average(self.min_var, batch_min,
                                             self.ema_decay)

        batch_max = self._reduce_list(tensors, tf.reduce_max)
        max_var = self.assign_moving_average(self.max_var, batch_max,
                                             self.ema_decay)
      else:
        min_var, max_var = self.min_var, self.max_var

      tensors = [
          tf.quantization.fake_quant_with_min_max_vars(
              tensor, min_var, max_var, num_bits=self.num_bits)
          for tensor in tensors
      ]
      tensor = tf.concat(tensors, axis=self.axis)
      return tf.quantization.fake_quant_with_min_max_vars(
          tensor, min_var, max_var, num_bits=self.num_bits)
    return tf.concat(tensors, axis=self.axis)

	# Copyright 2020 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.
	# ==============================================================================
	"""Layers for quantization."""

	import tensorflow as tf

	from layers import base_layers # import seq_flow_lite module


	class ActivationQuantization(base_layers.BaseLayer):
	"""A class that applies quantization to a activation tensor."""

	def __init__(self, ema_decay=0.99, num_bits=8, **kwargs):
	self.ema_decay = ema_decay
	self.num_bits = num_bits
	super(ActivationQuantization, self).__init__(**kwargs)

	def build(self, input_shapes):
	if self.parameters.quantize:
	self.min_var = self.add_weight(
	"min", initializer=tf.keras.initializers.Zeros(), trainable=False)
	self.max_var = self.add_weight(
	"max", initializer=tf.keras.initializers.Ones(), trainable=False)

	def call(self, inputs):
	if self.parameters.quantize:
	if self.parameters.mode == base_layers.TRAIN:
	# Toco expects 0.0 to be part of the quantization range.
	batch_min = tf.minimum(tf.reduce_min(inputs), 0.0)
	min_var = self.assign_moving_average(self.min_var, batch_min,
	self.ema_decay)

	batch_max = tf.maximum(tf.reduce_max(inputs), 0.0)
	max_var = self.assign_moving_average(self.max_var, batch_max,
	self.ema_decay)
	with tf.control_dependencies([min_var, max_var]):
	return tf.quantization.fake_quant_with_min_max_vars(
	inputs, batch_min, batch_max, num_bits=self.num_bits)
	else:
	return tf.quantization.fake_quant_with_min_max_vars(
	inputs, self.min_var, self.max_var, num_bits=self.num_bits)
	return inputs

	def quantize_using_range(self, inputs):
	# This method can only be called after a call to "call" method in this class
	if self.parameters.quantize:
	return tf.quantization.fake_quant_with_min_max_vars(
	inputs, self.min_var, self.max_var, num_bits=self.num_bits)
	return inputs


	class ConcatQuantization(ActivationQuantization):
	"""A class that applies quantization to a activation tensor."""

	def __init__(self, axis=2, **kwargs):
	self.axis = axis
	super(ConcatQuantization, self).__init__(**kwargs)

	def _reduce_list(self, tensor_list, functor):
	reduce_result = [functor(tensor) for tensor in tensor_list]
	# Toco expects 0.0 to be part of the quantization range.
	reduce_result.append(tf.constant(0.0))
	return functor(tf.stack(reduce_result))

	def call(self, tensors):
	# Ignore empty invocations done to build the keras layer.
	if tensors is None:
	return
	if self.parameters.quantize:
	if self.parameters.mode == base_layers.TRAIN:
	# Toco expects 0.0 to be part of the quantization range.
	batch_min = self._reduce_list(tensors, tf.reduce_min)
	min_var = self.assign_moving_average(self.min_var, batch_min,
	self.ema_decay)

	batch_max = self._reduce_list(tensors, tf.reduce_max)
	max_var = self.assign_moving_average(self.max_var, batch_max,
	self.ema_decay)
	else:
	min_var, max_var = self.min_var, self.max_var

	tensors = [
	tf.quantization.fake_quant_with_min_max_vars(
	tensor, min_var, max_var, num_bits=self.num_bits)
	for tensor in tensors
	]
	tensor = tf.concat(tensors, axis=self.axis)
	return tf.quantization.fake_quant_with_min_max_vars(
	tensor, min_var, max_var, num_bits=self.num_bits)
	return tf.concat(tensors, axis=self.axis)