# 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.

# ==============================================================================

"""Default 8-bit transforms."""

import collections

import inspect

import numpy as np

import tensorflow as tf

from tensorflow_model_optimization.python.core.quantization.keras import quantize_aware_activation

from tensorflow_model_optimization.python.core.quantization.keras import quantize_layer

from tensorflow_model_optimization.python.core.quantization.keras import quantizers

from tensorflow_model_optimization.python.core.quantization.keras import utils as quantize_utils

from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit import default_n_bit_quantize_configs as configs

from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit import default_n_bit_quantize_registry

from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations import transforms

try:

from keras.backend import unique_object_name # pylint: disable=g-import-not-at-top

except ImportError:

# Path as seen in pip packages as of TF/Keras 2.13.

from keras.src.backend import unique_object_name # pylint: disable=g-import-not-at-top

LayerNode = transforms.LayerNode

LayerPattern = transforms.LayerPattern

keras = tf.keras

def _get_conv_bn_layers(bn_layer_node):

bn_layer = bn_layer_node.layer

conv_layer = bn_layer_node.input_layers[0].layer

return conv_layer, bn_layer

def _get_weights(bn_layer_node):

+ list(bn_layer_node.weights.items()))

def _get_params(conv_layer, bn_layer, relu_layer=None):

return params

def _get_layer_node(fused_layer, weights):

return LayerNode(layer_config, weights, metadata=layer_metadata)

def _get_quantize_config(layer_node):

return layer_node.metadata.get('quantize_config')

def _has_custom_quantize_config(*layer_nodes):

return False

def _normalize_tuple(value):

return tuple(value)

class Conv2DBatchNormQuantize(transforms.Transform):

}

class Conv2DReshapeBatchNormQuantize(Conv2DBatchNormQuantize):

return self._replace(bn_layer_node, conv_layer_node)

class Conv2DBatchNormReLUQuantize(Conv2DBatchNormQuantize):

return self._replace(relu_layer_node, bn_layer_node, conv_layer_node)

class Conv2DBatchNormActivationQuantize(Conv2DBatchNormReLUQuantize):

inputs=[Conv2DBatchNormQuantize.pattern(self)])

class Conv2DReshapeBatchNormReLUQuantize(Conv2DBatchNormReLUQuantize):

return self._replace(relu_layer_node, bn_layer_node, conv_layer_node)

class Conv2DReshapeBatchNormActivationQuantize(

inputs=[Conv2DReshapeBatchNormQuantize.pattern(self)])

class DenseBatchNormQuantize(transforms.Transform):

}

class DenseBatchNormReLUQuantize(DenseBatchNormQuantize):

return self._replace(relu_layer_node, bn_layer_node, dense_layer_node)

class DenseBatchNormActivationQuantize(DenseBatchNormReLUQuantize):

inputs=[DenseBatchNormQuantize.pattern(self)])

class SeparableConv1DQuantize(transforms.Transform):

)])

class SeparableConvQuantize(transforms.Transform):

metadata=conv_metadata)

class LayerReLUQuantize(transforms.Transform):

}

class LayerReluActivationQuantize(LayerReLUQuantize):

inputs=[LayerPattern('Add|Conv2D|DepthwiseConv2D|Dense')])

class InputLayerQuantize(transforms.Transform):

}

class ConcatTransform(transforms.Transform):

# pylint:enable=protected-access

class ConcatTransform3Inputs(ConcatTransform):

inputs=[LayerPattern('.*'), LayerPattern('.*'), LayerPattern('.*')])

class ConcatTransform4Inputs(ConcatTransform):

LayerPattern('.*')])

class ConcatTransform5Inputs(ConcatTransform):

LayerPattern('.*'), LayerPattern('.*')])

class ConcatTransform6Inputs(ConcatTransform):

LayerPattern('.*'), LayerPattern('.*'), LayerPattern('.*')])