code: def main(unused_argv):
      ##############################################################################
      # Prepare training and testing data
      ##############################################################################
      (
          (x_train, y_train),
          (x_test, y_test),
          input_shape,
      ) = keras_test_utils.get_preprocessed_mnist_data()
  
      ##############################################################################
      # Train and convert a model with 2x2 block config. There's no kernel in tflite
      # supporting this block configuration, so the sparse tensor is densified and
      # the model falls back to dense execution.
      ##############################################################################
      pruning_params = {
          "pruning_schedule": ConstantSparsity(
              FLAGS.sparsity, begin_step=0, frequency=100
          ),
          "block_size": (2, 2),
      }
  
      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)
  
      # Get a dense model as baseline
      tflite_model_dense = converter.convert()
      tflite_model_path = "/tmp/dense_mnist.tflite"
      with open(tflite_model_path, "wb") as f:
          f.write(tflite_model_dense)
  
      # Enable sparse tensor encoding, otherwise the model is converted as dense.
      converter.optimizations = {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_2x2.tflite" % FLAGS.sparsity
      )
      with open(tflite_model_path, "wb") as f:
          f.write(tflite_model)
  
      print("evaluate 2x2 model")
      print(
          keras_test_utils.eval_mnist_tflite(model_content=tflite_model)
      )
  
      ##############################################################################
      # Train and convert a model with 1x4 block config. There's kernel in tflite
      # with this block configuration, so the model can take advantage of sparse
      # execution and see inference speed-up.
      ##############################################################################
      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 (4, 1) in the training API.
          "block_size": (4, 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.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_1x4.tflite" % FLAGS.sparsity
      )
      with open(tflite_model_path, "wb") as f:
          f.write(tflite_model)
  
      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)
      )
  
file_path: model-optimization/tensorflow_model_optimization/python/examples/sparsity/keras/mnist/mnist_e2e.py
imports: from __future__ import print_function
  from absl import app
  from absl import flags
  import tensorflow
  from tensorflow_model_optimization.python.core.keras import test_utils
  from tensorflow_model_optimization.python.core.sparsity.keras import prune
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_callbacks
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
label_desc: Use image node_id 3 for calling a global function with example usage: main(unused_argv) without return types