import tensorflow_model_optimization as tfmot

model = build_your_model()

pruning_params = {
  'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
                        initial_sparsity=0.0, final_sparsity=0.5,
                        begin_step=2000, end_step=4000)
}

pruned_model = tfmot.sparsity.keras.prune_low_magnitude(model, **pruning_params)