exporting .pt model into int8 quantized .engine model

[ausollet]

Hey, not sure if this has already been asked, but is it possible to export pyTorch .pt yolov5 models to tensorRT .engine int8 models.

[glenn-jocher]

@ausollet see export.py argparser for export arguments:

yolov5/export.py

Lines 479 to 508 in 6445a81

	def parse_opt():
	parser = argparse.ArgumentParser()
	parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
	parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
	parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
	parser.add_argument('--batch-size', type=int, default=1, help='batch size')
	parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
	parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
	parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
	parser.add_argument('--train', action='store_true', help='model.train() mode')
	parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
	parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
	parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
	parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
	parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
	parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
	parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
	parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
	parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model')
	parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
	parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
	parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
	parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
	parser.add_argument('--include', nargs='+',
	default=['torchscript', 'onnx'],
	help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs')
	opt = parser.parse_args()
	print_args(FILE.stem, opt)
	return opt

[ausollet]

Just changing the config flag to int8 gave rise to errors.
Am I missing something? Do I need to define a calibrator too for int8 conversion?

[glenn-jocher]

@ausollet the only available TRT arguments are:

[glenn-jocher]

@ausollet --half is also compatible with TRT export:

yolov5/export.py

Line 177 in 4c40933

def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):

[ausollet]

Ok, so basically from what I understand, as of now in export.py TRT .engine can only be either FP32 or FP16 (half).
int8 support doesn't exist yet.
Point me out of I'm incorrect.

[glenn-jocher]

@ausollet yes that's correct! I'm not a TRT expert so if there's a way to do this please help by submitting a PR.

Please see our ✅ Contributing Guide to get started.