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 323 forks source link

Model Pruning with Yolo Object Detection Model #1018

Open divyajaincs opened 2 years ago

divyajaincs commented 2 years ago

Is there any way to use tfmot for Keras Yolov3-v4 model pruning, or any other object detection model.

rino20 commented 2 years ago

Hi Divya,

Sorry, there is not yet an official repository using Pruning to object detection models. But of course, you can try one with Tensorflow Model garden models.

There are some object detection models w/ RetinaNet, https://github.com/tensorflow/models/tree/master/official/vision/configs/experiments/retinanet

And you can apply the pruning similarly in image classification examples. https://github.com/tensorflow/models/blob/master/official/projects/pruning/tasks/image_classification.py

Note that, there are some subclass structures of usual object detection model, so you may need to deal with that manually like this - https://github.com/tensorflow/models/blob/master/official/projects/pruning/tasks/image_classification.py#L130-L133

Let us know your target model code, if you need more help.