Pruning not working for tf.keras.Batchnorm

[jkparuchuri]

Describe the bug
ValueError: Please initialize Prune with a supported layer. Layers should either be a PrunableLayer instance, or should be supported by the PruneRegistry. You passed: <class 'tensorflow.python.keras.layers.normalization.BatchNormalization'>

System information

TensorFlow installed from (source or binary): binary

TensorFlow version: 2.1.0

TensorFlow Model Optimization version: 0.2.1

Python version: 3.5.6

[alanchiao]

@jkparuchuri : how are you creating your batchnorm layer? In TF 2.1.0, the batchnorm layer class should be tensorflow.python.keras.layers.normalization_v2.BatchNormalization if you are using the tf.keras.layers.BatchNormalization API.

If you are using tf.compat.v1.keras.layers.BatchNormalization in 2.1.0, that is currently not supported. I can update the documentation in that case.

[alanchiao]

Making a documentation update in #225.

[jkparuchuri]

@alanchiao I didnt use compat.v1 but the keras model was saved using .h5 format. Now I tried to save in in default saved model format and loaded for it. I gave a different error aa below. Any suggestions.

Please initialize Prune with a supported layer. Layers should either be a PrunableLayer instance, or should be supported by the PruneRegistry. You passed: <class 'tensorflow.python.keras.saving.saved_model.load.ZeroPadding2D'>

[alanchiao]

Thanks @jkparuchuri. Looks like the keras model was originally built in TF 1.X and saved using h5 and you are now reloading it in TF 2.X. This isn't a case I had considered actually and glad you brought it up. The same with regards to SavedModel (it uses different classes under the hood right now when compared to using the tf.keras.layers directly).

In the mean time, a workaround (without needing the original model architecture) would be to modify this file and add a row below "layers.BatchNormalization: []" that says "tf.compat.v1.layers.BatchNormalization: []". Then build the pip package via https://github.com/tensorflow/model-optimization#installing-from-source. Keep using the h5 model in that case.

If you're using one of the common vision models (e.g. mobilenet), that should suffice.

I'll consider what is appropriate for your case.

[jkparuchuri]

@alanchiao Building from source not working on my mac

ITs giving error ERROR: no such target '//tensorflow_model_optimization:pip_pkg': target 'pip_pkg' not declared in package 'tensorflow_model_optimization' defined by /Users/jith/june/git/model-optimization/tensorflow_model_optimization/BUILD.
INFO: Elapsed time: 0.083s

[alanchiao]

Can you cd out of tensorflow_model_optimization into the model-optimization folder? Then run the same command.

I realized that the cd tensorflow_model_optimization instruction is incorrect. Should say cd model-optimization.

[jkparuchuri]

@alanchiao Tried that but its giving different errors

ERROR: /Users/jith/june/git/model-optimization/tensorflow_model_optimization/python/core/clustering/keras/BUILD:63:1: //tensorflow_model_optimization/python/core/clustering/keras:cluster_test: no such attribute 'python_version' in 'py_test' rule.
ERROR: /Users/jith/june/git/model-optimization/tensorflow_model_optimization/python/core/clustering/keras/BUILD:75:1: //tensorflow_model_optimization/python/core/clustering/keras:clustering_centroids_test: no such attribute 'python_version' in 'py_test' rule.
ERROR: /Users/jith/june/git/model-optimization/tensorflow_model_optimization/python/core/clustering/keras/BUILD:87:1: //tensorflow_model_optimization/python/core/clustering/keras:cluster_wrapper_test: no such attribute 'python_version' in 'py_test' rule.
ERROR: /Users/jith/june/git/model-optimization/tensorflow_model_optimization/python/core/clustering/keras/BUILD:103:1: //tensorflow_model_optimization/python/core/clustering/keras:clustering_registry_test: no such attribute 'python_version' in 'py_test' rule.
ERROR: /Users/jith/june/git/model-optimization/tensorflow_model_optimization/python/core/clustering/keras/BUILD:9:1: Target '//tensorflow_model_optimization/python/core/clustering/keras:cluster' contains an error and its package is in error and referenced by '//tensorflow_model_optimization/python/core/clustering/keras:keras'.
ERROR: /Users/jith/june/git/model-optimization/tensorflow_model_optimization/python/core/clustering/keras/BUILD:9:1: Target '//tensorflow_model_optimization/python/core/clustering/keras:init.py' contains an error and its package is in error and referenced by '//tensorflow_model_optimization/python/core/clustering/keras:keras'.
ERROR: /Users/jith/june/git/model-optimization/tensorflow_model_optimization/python/core/clustering/BUILD:5:1: Target '//tensorflow_model_optimization/python/core/clustering/keras:keras' contains an error and its package is in error and referenced by '//tensorflow_model_optimization/python/core/clustering:clustering'.
ERROR: Analysis of target '//:pip_pkg' failed; build aborted: Loading failed.
INFO: Elapsed time: 0.136s

[alanchiao]

python_version is available in bazel since 0.22.0 (see https://docs.bazel.build/versions/0.22.0/be/python.html). Do you think you'd be able to upgrade your version of bazel? I'll need to take a look at the TFMOT docs and clarify the minimum bazel version also.

[jkparuchuri]

@alanchiao Thankyou, upgrading bazel resolved the building from source. But still see the same error even after adding to registry as you suggested :(

'PruneRegistry. You passed: {input}'.format(input=layer.class))
ValueError: Please initialize Prune with a supported layer. Layers should either be a PrunableLayer instance, or should be supported by the PruneRegistry. You passed: <class 'tensorflow.python.keras.layers.normalization.BatchNormalization'>

[alanchiao]

I'd make sure that the tensorflow_model_optimization you're using is the one built from https://github.com/tensorflow/model-optimization#installing-from-source with your change as opposed to the one you previously installed that MOT provided. I'm pretty confident what I suggested was correct.

[jkparuchuri]

@alanchiao Able to resolve the issue. Its "tf.compat.v1.keras.layers.BatchNormalization: []". keras was missing earlier in that line.

[bilalahmed381]

Trying to prune Resnet50 and getting this error, can you tell me what't the solution of this.

ValueError: Please initialize Prune with a supported layer. Layers should either be a PrunableLayer instance, or should be supported by the PruneRegistry. You passed: <class 'tensorflow.python.keras.layers.normalization.BatchNormalization'>

[jkparuchuri]

@Craftsman381 Workaround is to add "tf.compat.v1.keras.layers.BatchNormalization: []" to prune_registry at https://github.com/tensorflow/model-optimization/blob/master/tensorflow_model_optimization/python/core/sparsity/keras/prune_registry.py#L82 and rebuild from source