def apply(self, model, layer_quantize_map):
"""Implement default 8-bit transforms.
Currently this means the following.
1. Pull activations into layers, and apply fuse activations. (TODO)
2. Modify range in incoming layers for Concat. (TODO)
3. Fuse Conv2D/DepthwiseConv2D + BN into single layer.
Args:
model: Keras model to be quantized.
layer_quantize_map: Map with keys as layer names, and values as dicts
containing custom `QuantizeConfig`s which may have been passed with
layers.
Returns:
(Transformed Keras model to better match TensorFlow Lite backend, updated
layer quantize map.)
"""
transforms = [
default_8bit_transforms.InputLayerQuantize(),
default_8bit_transforms.Conv2DBatchNormReLUQuantize(),
default_8bit_transforms.Conv2DBatchNormActivationQuantize(),
default_8bit_transforms.Conv2DBatchNormQuantize(),
default_8bit_transforms.ConcatTransform6Inputs(),
default_8bit_transforms.ConcatTransform5Inputs(),
default_8bit_transforms.ConcatTransform4Inputs(),
default_8bit_transforms.ConcatTransform3Inputs(),
default_8bit_transforms.ConcatTransform(),
default_8bit_transforms.AddReLUQuantize(),
default_8bit_transforms.AddActivationQuantize(),
]
return model_transformer.ModelTransformer(
model, transforms, layer_quantize_map.keys(),
layer_quantize_map).transform()
def testAddsQuantizeLayerAfterInputLayer(self):
inp1 = keras.layers.Input((3, ))
inp2 = keras.layers.Input((3, ))
x = keras.layers.Concatenate()([inp1, inp2])
model = keras.Model([inp1, inp2], x)
transformed_model, _ = ModelTransformer(
model, [default_8bit_transforms.InputLayerQuantize()]).transform()
for input_layer in transformed_model._input_layers:
layer_after_input = input_layer._outbound_nodes[0].outbound_layer
self.assertIsInstance(layer_after_input,
quantize_layer.QuantizeLayer)
self.assertIsInstance(layer_after_input.quantizer,
quantizers.MovingAverageQuantizer)
