import tensorflow as tf import tensorflow_model_optimization as tfmot model = tf.keras.Sequential([...]) pruning_schedule = tfmot.sparsity.keras.PolynomialDecay( initial_sparsity=0.0, final_sparsity=0.5, begin_step=2000, end_step=4000) model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude( model, pruning_schedule=pruning_schedule) ... model_for_pruning.fit(...) TensorFlow Model Optimization