tensorflow / tensorflow/model-optimization
Unable to prune/quantize multiple layers at the same time
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.6k
- Forks
- 349
- Avg merge
- 3d 2h
- Merged PRs (30d)
- 1
Description
Describe the bug
Unable to prune/quantize multiple layers at the same time
System information
TensorFlow version (installed from source or binary): 2.4.0
TensorFlow Model Optimization version (installed from source or binary): 0.7.2
Python version: 3.6.9
Describe the expected behavior
Ability to prune/quantize multiple layers like Conv2D and Dense layers at the same time.
Describe the current behavior
Current API supports pruning/quantization of either Conv2D or Dense layers at a time, not both.
Code to reproduce the issue
Provide a reproducible code that is the bare minimum necessary to generate the
problem.
For Pruning - Only Conv2D layers(Don't be confused, problem is inability to combine the both)
Code Ref(Modified) - https://www.tensorflow.org/model_optimization/guide/pruning/comprehensive_guide
# Create a base model
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
# Helper function uses `prune_low_magnitude` to make only the
# Dense layers train with pruning.
def apply_pruning_to_conv2d(layer):
if isinstance(layer, tf.keras.layers.Conv2D):
return tfmot.sparsity.keras.prune_low_magnitude(layer)
return layer
# Use `tf.keras.models.clone_model` to apply `apply_pruning_to_dense`
# to the layers of the model.
model_for_pruning = tf.keras.models.clone_model(
base_model,
clone_function=apply_pruning_to_conv2d,
)
model_for_pruning.summary()
For Quantization - Only Dense layers(Don't be confused, problem is inability to combine the both)
Code Ref - https://www.tensorflow.org/model_optimization/guide/quantization/training_comprehensive_guide
# Create a base model
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
# Helper function uses `quantize_annotate_layer` to annotate that only the
# Dense layers should be quantized.
def apply_quantization_to_dense(layer):
if isinstance(layer, tf.keras.layers.Dense):
return tfmot.quantization.keras.quantize_annotate_layer(layer)
return layer
# Use `tf.keras.models.clone_model` to apply `apply_quantization_to_dense`
# to the layers of the model.
annotated_model = tf.keras.models.clone_model(
base_model,
clone_function=apply_quantization_to_dense,
)
# Now that the Dense layers are annotated,
# `quantize_apply` actually makes the model quantization aware.
quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model)
quant_aware_model.summary()
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
No source files or tests are named. Start by reproducing the pruning and quantization examples using clone_model, prune_low_magnitude, quantize_annotate_layer, and quantize_apply; done means both Conv2D and Dense layers can be handled together through the supported API.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- keras, python, tensorflow
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 28/100