tensorflow / tensorflow/model-optimization

Unable to prune/quantize multiple layers at the same time

Open
#955 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.