NVIDIA-AI-IOT / yolo_deepstream

yolo model qat and deploy with deepstream&tensorrt
Apache License 2.0
533 stars 135 forks source link

Add support for quantization with custom datasets #30

Open amansrf opened 1 year ago

amansrf commented 1 year ago

Modify qat.py to allow custom dataset quantization. Note: Benchmark on coco will not work when using custom dataset.

Tested on: Ubuntu 20.04 RTX 3090 CUDA 11.6 TensorRT 8.4.3-1+cuda11.6 PyTorch 1.12.1+cu116