From e84e11724b746ff577f28f332c231e429dd180c0 Mon Sep 17 00:00:00 2001 From: sifan-intel <109498696+sifan-intel@users.noreply.github.com> Date: Thu, 28 Mar 2024 04:00:57 +0800 Subject: [PATCH] Specify YOLODataset task arg during INT8 calibration (#9309) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- ultralytics/engine/exporter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 859106a1..b73ecf5c 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -459,7 +459,7 @@ class Exporter: # Generate calibration data for integer quantization LOGGER.info(f"{prefix} collecting INT8 calibration images from 'data={self.args.data}'") data = check_det_dataset(self.args.data) - dataset = YOLODataset(data["val"], data=data, imgsz=self.imgsz[0], augment=False) + dataset = YOLODataset(data["val"], data=data, task=self.model.task, imgsz=self.imgsz[0], augment=False) n = len(dataset) if n < 300: LOGGER.warning(f"{prefix} WARNING ⚠️ >300 images recommended for INT8 calibration, found {n} images.")