Add an example of combining sparsity & dynamic range quantization.

PiperOrigin-RevId: 361198975

Author: liyunlu0618

tensorflow_model_optimization/python/examples/sparsity/keras/mnist/mnist_e2e.py

@@ -164,6 +164,37 @@ def main(unused_argv):
   print('evaluate 1x4 model')
   print(keras_test_utils.eval_mnist_tflite(model_content=tflite_model))
 
+  ##############################################################################
+  # Train and convert a model with 1x16 block config, and enable post-training
+  # dynamic range quantization during conversion.
+  ##############################################################################
+  pruning_params = {
+      'pruning_schedule':
+          ConstantSparsity(FLAGS.sparsity, begin_step=0, frequency=100),
+      # TFLite transposes the weight during conversion, so we need to specify
+      # the block as (16, 1) in the training API.
+      'block_size': (16, 1)
+  }
+
+  model = build_layerwise_model(input_shape, **pruning_params)
+  model = train(model, x_train, y_train, x_test, y_test)
+
+  converter = tf.lite.TFLiteConverter.from_keras_model(model)
+  converter.optimizations = {
+      tf.lite.Optimize.DEFAULT, tf.lite.Optimize.EXPERIMENTAL_SPARSITY
+  }
+
+  tflite_model = converter.convert()
+  # Check the model is compressed
+  print('Compression ratio: ', len(tflite_model) / len(tflite_model_dense))
+
+  tflite_model_path = '/tmp/sparse_mnist_%s_1x16.tflite' % FLAGS.sparsity
+  with open(tflite_model_path, 'wb') as f:
+    f.write(tflite_model)
+
+  print('evaluate 1x16 model')
+  print(keras_test_utils.eval_mnist_tflite(model_content=tflite_model))
+
 
 if __name__ == '__main__':
   absl_app.run(main)