Tensorrt-int8-quantization-pipline

a simple pipline of int8 quantization based on tensorrt.
classification detection

Example for classification

cd classification

1、Choose a model and prepare a calibration dataset，like resnet101 training from imagenet1k.

wget https://hanlab.mit.edu/files/OnceForAll/ofa_cvpr_tutorial/imagenet_1k.zip
unzip 'imagenet_1k.zip'
mkdir model

2、eval the float32 model performance.

python test_torch.py

3、convert to onnx model.

python torch2onnx.py

4、 quantization int8 trt model.

python quantization.py

5、eval the int8 model performance.

python test_int8trt.py

or run a pipline including the above steps.

python tensorrt_PTA_classification_pipline.py

Note

You can replace resnet101 with your network. If your dataset structure is different, you need to modify some code about dataset.

# test_torch.py torch2onnx.py quantization.py
if __name__ == "__main__":
    net = models.resnet101(pretrained=True).to('cpu')

or

# tensorrt_PTA_classification_pipline.py
if __name__ == "__main__":
    net = models.resnet101(pretrained=True).to('cpu')

Example for detection

cd detection

1、Choose a model and test inference，like YOLOX-s.

wget wget https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_s.onnx
python demo_onnx.py --model_path yolox_s.onnx --label_name_path coco.label --image_path dog.jpg --output_path output_onnx.jpg

2、random sample 2k training images as calibration data, YOLOX-s training from COCO2017.

mkdir calibration
python sample.py --traing_data_path your_path/coco/images/train2017/  --count 2000 --calibration_path ./calibration/

3、quantization

python3 -m onnxsim yolox_s.onnx yolox_s.onnx
python quantization.py

4、test int tensort model

python demo_trt.py --model_path modelInt8.engine --label_name_path coco.label --image_path dog.jpg --output_path output_trt.jpg

float32 onnx	int8 tensorrt

TO DO

• [x] example for detection.
• [ ] example for segmention.