Files
wehub-resource-sync c8a779b1bb
Docker Image CI / build-ubuntu2004 (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:55 +08:00

104 lines
5.4 KiB
Markdown
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# About
This folder contains the Quantization-Aware Training (QAT) workflow for [standard networks](#step-1-model-quantization-and-fine-tuning).
The QAT end-to-end workflow (TF2-to-ONNX) consists of the following steps:
- Model quantization using the `quantize_model` function with `NVIDIA` quantization scheme.
- QAT model fine-tuning (saves checkpoints).
- Baseline vs QAT models accuracy comparison.
- QAT model conversion to SavedModel format.
- Conversion of SavedModel to ONNX.
- TensorRT engine building via ONNX file and inference.
# Requirements
## 1. Base requirements
1. Install `tensorflow-quantization` toolkit.
2. Install additional requirements: `pip install -r requirements.txt`.
3. (Optional) Install TensorRT for full workflow support (needed for `infer_engine.py`).
**Note**: For CLI run, please go to the cloned repository's root directory and run `export PYTHONPATH=$PWD`, so that the `examples` folder is available for import.
## 2. Data preparation
### A. Raw data download
We are using the ImageNet 2012 dataset (task 1 - image classification), which requires manual downloads due to terms of access agreements.
Please login/sign-up on [the ImageNet website](https://image-net.org/challenges/LSVRC/2012/2012-downloads.php) and download the "train/validation data".
This is needed for the QAT model fine-tuning, and it is also used to evaluate the Baseline and QAT models.
### B. Conversion to tfrecord
Our workflow supports `tfrecord` format, so please follow the following instructions (modified from [TensorFlow's instructions](https://github.com/tensorflow/tpu/tree/master/tools/datasets#imagenet_to_gcspy)) to convert the downloaded `.tar` ImageNet files to the required format:
1. Set `IMAGENET_HOME=/path/to/imagenet/tar/files` in [`data/imagenet_data_setup.sh`](data/imagenet_data_setup.sh).
2. Download [`imagenet_to_gcs.py`](https://github.com/tensorflow/tpu/blob/master/tools/datasets/imagenet_to_gcs.py) to `$IMAGENET_HOME`.
3. Run `./data/imagenet_data_setup.sh`.
# Workflow
## Step 1: Model quantization and fine-tuning
Model quantization, fine-tuning, and conversion to ONNX.
Example models:
| Model | Task | Script - QAT Workflow |
|---------------|------------------|------------------------------|
| ResNet | Classification | [resnet](resnet) |
| EfficientNet | Classification | [efficientnet](efficientnet) |
| MobileNet | Classification | [mobilenet](mobilenet) |
| Inception | Classification | [inception](inception) |
> For each model's performance results, please refer to the toolkit's User Guide ("Model Zoo").
## Step 2: TensorRT deployment
Build the TensorRT engine and evaluate its latency and accuracy performances.
#### 2.1. Build TensorRT engine from ONNX
Convert the ONNX model into a TensorRT engine (also obtains latency measurements):
```sh
trtexec --onnx=model_qat.onnx --int8 --saveEngine=model_qat.engine --verbose
```
Arguments:
* `--onnx`: Path to QAT onnx graph.
* `--saveEngine`: Output filename of TensorRT engine.
* `--verbose`: Flag to enable verbose logging.
#### 2.2. TensorRT Inference
Obtain accuracy results on the validation dataset:
```sh
python infer_engine.py --engine=<path_to_trt_engine> --data_dir=<path_to_tfrecord_val_data> -b=<batch_size>
```
Arguments:
- `-e, --engine`: TensorRT engine filename (to load).
- `-m, --model_name`: Name of the model, needed to choose the appropriate input pre-processing. Options={`resnet_v1` (default), `resnet_v2`, `efficientnet_b0`, `efficientnet_b3`, `mobilenet_v1`, `mobilenet_v2`}.
- `-d, --data_dir`: Path to directory of input images in **tfrecord format** (`data["validation"]`).
- `-k, --top_k_value` (default=1): Value of `K` for the top-K predictions used in the accuracy calculation.
- `-b, --batch_size` (default=1): Number of inputs to send in parallel (up to max batch size of engine).
- `--log_file`: Filename to save logs.
Outputs:
- `.log` file: contains the engine's performance accuracy.
# Additional resources
The following resources provide a deeper understanding about Quantization aware training, TF2ONNX and importing a model into TensorRT using Python.
**Quantization Aware Training**
* <a href="https://developer.nvidia.com/blog/achieving-fp32-accuracy-for-int8-inference-using-quantization-aware-training-with-tensorrt/">Achieving FP32 Accuracy for INT8 Inference Using Quantization Aware Training with NVIDIA TensorRT</a>
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://arxiv.org/pdf/1712.05877.pdf)
- [Quantization Aware Training guide](https://www.tensorflow.org/model_optimization/guide/quantization/training)
- [Deep Residual Learning for Image Recognition](https://arxiv.org/pdf/1512.03385.pdf)
**Parsers**
- [TF2ONNX Converter](https://github.com/onnx/tensorflow-onnx)
- [ONNX Parser](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/parsers/Onnx/pyOnnx.html)
**Documentation**
- [Introduction To NVIDIAs TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
- [Importing A Model Using A Parser In Python](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#import_model_python)
- [NVIDIAs TensorRT Documentation](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html)