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)