tensorflow / model-optimization

A toolkit to optimize ML models for deployment for Keras and TensorFlow, including quantization and pruning.
https://www.tensorflow.org/model_optimization
Apache License 2.0
1.49k stars 319 forks source link

Can't run prunning with XNNPack #837

Open jhelsas opened 3 years ago

jhelsas commented 3 years ago

I am having the following error trying to prune an Efficientnet B1 network, but I am unable to understand even what it is referring to:

Could not find `Conv2D 3x3` layer with stride 2x2, `input filters == 3` and `VALID` padding and preceding `ZeroPadding2D` with `padding == 1` in all input branches of the model

I am calling tfmot.sparsity.keras.prune_low_magnitude with the following args:

    pruning_params = {
        'pruning_schedule' : tfmot.sparsity.keras.PolynomialDecay(
                                              initial_sparsity=0.25,
                                              final_sparsity=0.75,
                                              begin_step=0,
                                              end_step=end_step),
        'pruning_policy' : tfmot.sparsity.keras.PruneForLatencyOnXNNPack()
        }

As described here: https://www.tensorflow.org/model_optimization/guide/pruning/pruning_for_on_device_inference

My efficientnet doesn't have any ZeroPadding2D layer, so I don't get why tfmot.sparsity.keras.prune_low_magnitude is throwing that error.

I do have some conv2d with 2x2 strides:

conv2d layer
name:  stem_conv
filters:  32
kernel_size:  (3, 3)
strides:  (2, 2)
padding:  same

conv2d layer
name:  block2a_dwconv
filters:  None
kernel_size:  (3, 3)
strides:  (2, 2)
padding:  same

conv2d layer
name:  block3a_dwconv
filters:  None
kernel_size:  (5, 5)
strides:  (2, 2)
padding:  same

conv2d layer
name:  block4a_dwconv
filters:  None
kernel_size:  (3, 3)
strides:  (2, 2)
padding:  same

conv2d layer
name:  block6a_dwconv
filters:  None
kernel_size:  (5, 5)
strides:  (2, 2)
padding:  same

But there is no related ZeroPadding2D layer anywhere, much less before any of those. Changing the padding type to "valid" don't solve the issue.

I am really lost on how to solve this, and any help is welcome. Thanks in advance.

rino20 commented 2 years ago

Hi @jhelsas, Thanks for reporting this.

The Pruning policy you are using (PruneForLatencyOnXNNPack) is customized for a specific subgraph structure as shown in the guide's example. It searches for a specific conv2d 3x3 layer after Zeropadding2D, so the error message was generated regardless of your model.

Just remove the pruning policy, and follow the further steps in the guide document. I think you can still get a decent performance for the efficient net without the policy.

Please let me know if you need any assistance.