README.md

# Training with Pruning
[TOC]

⚠️ Disclaimer: All datasets hyperlinked from this page are not owned or
distributed by Google. The dataset is made available by third parties.
Please review the terms and conditions made available by the third parties
before using the data.

## Overview

This project includes pruning codes for TensorFlow models.
These are examples to show how to apply the Model Optimization Toolkit's
[pruning API](https://www.tensorflow.org/model_optimization/guide/pruning).

## How to train a model

```bash
EXPERIMENT=xxx  # Change this for your run, for example, 'resnet_imagenet_pruning'
CONFIG_FILE=xxx  # Change this for your run, for example, path of imagenet_resnet50_pruning_gpu.yaml
MODEL_DIR=xxx  #  Change this for your run, for example, /tmp/model_dir
python3 train.py \
  --experiment=${EXPERIMENT} \
  --config_file=${CONFIG_FILE} \
  --model_dir=${MODEL_DIR} \
  --mode=train_and_eval
```

## Accuracy
<figure align="center">
<img width=70% src=https://storage.googleapis.com/tf_model_garden/models/pruning/images/readme-pruning-classification-resnet.png>
<img width=70% src=https://storage.googleapis.com/tf_model_garden/models/pruning/images/readme-pruning-classification-mobilenet.png>
<figcaption>Comparison of Imagenet top-1 accuracy for the classification models</figcaption>
</figure>

Note: The Top-1 model accuracy is measured on the validation set of [ImageNet](https://www.image-net.org/).

## Pre-trained Models

### Image Classification

Model |Resolution|Top-1 Accuracy (Dense)|Top-1 Accuracy (50% sparsity)|Top-1 Accuracy (80% sparsity)|Config |Download
----------------------|----------|---------------------|-------------------------|-------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------|
|MobileNetV2           |224x224   |72.768%              |71.334%                  |61.378%                 |[config](https://github.com/tensorflow/models/blob/master/official/projects/pruning/configs/experiments/image_classification/imagenet_mobilenetv2_pruning_gpu.yaml)  |[TFLite(50% sparsity)](https://storage.googleapis.com/tf_model_garden/vision/mobilenet/v2_1.0_float/mobilenet_v2_0.5_pruned_1.00_224_float.tflite),                   |
|ResNet50              |224x224   |76.704%              |76.61%                  |75.508%                 |[config](https://github.com/tensorflow/models/blob/master/official/projects/pruning/configs/experiments/image_classification/imagenet_resnet50_pruning_gpu.yaml)     |[TFLite(80% sparsity)](https://storage.googleapis.com/tf_model_garden/vision/resnet50_imagenet/resnet_50_0.8_pruned_224_float.tflite)                                |