def apply(self, model, layer_quantize_map):
"""Implement TFLite 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 `QuantizeProvider`s which may have been passed with
layers.
Returns:
Transformed Keras model to better match TensorFlow Lite backend.
"""
# TODO(pulkitb): Sequential models not supported yet. Remove once support is
# added.
if isinstance(model, keras.Sequential):
return model
transforms = [
tflite_transforms.DepthwiseConv2DBatchNormReLU6Fold(),
tflite_transforms.Conv2DBatchNormReLU6Fold(),
tflite_transforms.Conv2DBatchNormFold(),
]
return model_transformer.ModelTransformer(
model, transforms, layer_quantize_map.keys(),
layer_quantize_map).transform()
def testTransformsDepthwiseConvBNReLUPattern(self):
model = DepthwiseConv2DModel.get_nonfolded_batchnorm_model(
post_bn_activation=keras.layers.ReLU(6.0), model_type='functional')
folded_model = DepthwiseConv2DModel.get_folded_batchnorm_model(
post_bn_activation=keras.layers.ReLU(6.0), is_quantized=True)
transformed_model, _ = ModelTransformer(
model, [tflite_transforms.DepthwiseConv2DBatchNormReLU6Fold()
]).transform()
inputs = np.random.standard_normal(
DepthwiseConv2DModel.get_batched_input_shape())
self.assertAllClose(transformed_model.predict(inputs),
folded_model.predict(inputs))
def testTransformsDepthwiseConvBNReLUPatternPreservesWeights(self):
# random_init to prevent non-random initialization in resulting
# in same weights between transformed and non-transformed models.
model = DepthwiseConv2DModel.get_nonfolded_batchnorm_model(
post_bn_activation=keras.layers.ReLU(6.0),
model_type='functional',
random_init=True)
transformed_model, _ = ModelTransformer(
model, [tflite_transforms.DepthwiseConv2DBatchNormReLU6Fold()
]).transform()
transformed_weights = transformed_model.get_weights()
# Remove quantization related weights.
del transformed_weights[3:8]
self.assertEqual(len(transformed_weights), len(model.get_weights()))
for i in range(len(transformed_weights)):
self.assertAllEqual(transformed_weights[i], model.get_weights()[i])
