NVIDIA-AI-IOT / yolo_deepstream

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

Add support for quantization with custom datasets #30

Open amansrf opened 2 years ago

amansrf commented 2 years 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