wehub-resource-sync
103 lines
4.1 KiB
Markdown
103 lines
4.1 KiB
Markdown
## About
|
|
This script presents a QAT end-to-end workflow (TF2-to-ONNX) for [EfficientNet](https://github.com/tensorflow/models/tree/master/official/legacy/image_classification/efficientnet).
|
|
|
|
### Contents
|
|
[Requirements](#requirements) • [Workflow](#workflow) • [Results](#results)
|
|
|
|
## Requirements
|
|
1. Install base requirements and prepare data. Please refer to [examples' README](../README.md).
|
|
|
|
2. Clone the models from Tensorflow model garden:
|
|
|
|
```
|
|
git clone https://github.com/tensorflow/models.git
|
|
pushd models && git checkout tags/v2.8.0 && popd
|
|
export PYTHONPATH=$PWD/models:$PYTHONPATH
|
|
pip install -r models/official/requirements.txt
|
|
```
|
|
|
|
> cd models && git submodule init && git submodule update
|
|
|
|
3. Download pretrained checkpoints:
|
|
1. B0: https://tfhub.dev/tensorflow/efficientnet/b0/classification/1
|
|
2. B3: https://tfhub.dev/tensorflow/efficientnet/b3/classification/1
|
|
|
|
|
|
## Workflow
|
|
|
|
### Step 1: Model Quantization and Fine-tuning
|
|
|
|
* In `run_qat_workflow.py`, please set the `pretrained_ckpt_path` field to the directory of the downloaded checkpoint to start fine-tuning with QAT. All the required hyper-parameters can be set in the `HYPERPARAMS` dictionary.
|
|
|
|
Please run the following to quantize, fine-tune, and save the final graph in SavedModel format.
|
|
|
|
```sh
|
|
python run_qat_workflow.py
|
|
```
|
|
> Update `MODEL_VERSION` to the EfficientNet version you wish to quantize.
|
|
|
|
### Step 2: Exporting a QAT SavedModel
|
|
|
|
Once you've fine-tuned the QAT model, export it by running
|
|
|
|
```sh
|
|
python export.py --ckpt <path_to_pretrained_ckpt> --output <saved_model_output_name> --model_version b0
|
|
```
|
|
|
|
This script applies quantization to the model, restores the checkpoint, and exports it in a SavedModel format. This script will generate `eff` which is a directory containing saved model. We set the overall graph data format to `NCHW` by using `tf.keras.backend.set_image_data_format('channels_first')`. TensorRT expects `NCHW` format for graphs trained with QAT for better optimizations.
|
|
|
|
Arguments:
|
|
|
|
* `--ckpt` : Path to fine-tuned QAT checkpoint to be loaded.
|
|
* `--output` : Name of output TF saved model.
|
|
* `--model_version` : EfficientNet model version, currently supports {`b0`, `b3`}.
|
|
|
|
### Step 3: Conversion to ONNX
|
|
|
|
Convert the saved model into ONNX by running
|
|
|
|
```sh
|
|
python -m tf2onnx.convert --saved-model <path_to_saved_model> --output model_qat.onnx --opset 13
|
|
```
|
|
|
|
By default, tf2onnx uses TF's graph optimizers to performs constant folding after a saved model is loaded.
|
|
|
|
Arguments:
|
|
|
|
* `--saved-model` : Name of TF SavedModel
|
|
* `--output` : Name of ONNX output graph
|
|
* `--opset` : ONNX opset version (opset 13 or higher must be used)
|
|
|
|
### Step 4: TensorRT Deployment
|
|
Please refer to the [examples' README](../README.md).
|
|
|
|
## Results
|
|
|
|
This section presents the validation accuracy for the full ImageNet dataset on NVIDIA's A100 GPU and TensorRT 8.4 GA.
|
|
|
|
### EfficientNet-B0
|
|
| Model | Accuracy (%) | Latency (ms, bs=1) |
|
|
|-----------------------|--------------|--------------------|
|
|
| Baseline (TensorFlow) | 76.97 | 6.77 |
|
|
| PTQ (TensorRT) | 71.71 | 0.67 |
|
|
| **QAT** (TensorRT) | 75.82 | 0.68 |
|
|
> QAT fine-tuning hyper-parameters: `bs64, ep10, lr=0.001, steps_per_epoch=None`
|
|
|
|
### EfficientNet-B3
|
|
| Model | Accuracy (%) | Latency (ms, bs=1) |
|
|
|-----------------------|--------------|--------------------|
|
|
| Baseline (TensorFlow) | 81.36 | 10.33 |
|
|
| PTQ (TensorRT) | 78.88 | 1.24 |
|
|
| **QAT** (TensorRT) | 79.48 | 1.23 |
|
|
> QAT fine-tuning hyper-parameters: `bs=32, ep20, steps_per_epoch=None, lr=0.0001`
|
|
|
|
### Notes
|
|
- QAT fine-tuning hyper-parameters:
|
|
- Optimizer: `piecewise_sgd`, `lr_schedule=[(1.0, 1), (0.1, 3), (0.01, 6), (0.001, 9), (0.001, 15)]`.
|
|
- Other hyper-parameters are under each model's results table.
|
|
- PTQ calibration: `bs=64`
|
|
- EfficientNet model quantization:
|
|
- QDQ nodes added in Residual connection (fix added to ResidualQDQCustomCase for `Conv-BN-Activation-Dropout` pattern),
|
|
- Global Average Pooling,
|
|
- Multiply layer in SE block.
|