Fix TFLite INT8 quant bug (#13082)

.github/workflows/ci.yaml

@@ -164,7 +164,7 @@ jobs:
 
   Tests:
     if: github.event_name != 'workflow_dispatch' || github.event.inputs.tests == 'true'
-    timeout-minutes: 60
+    timeout-minutes: 120
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -241,7 +241,7 @@ jobs:
 
   RaspberryPi:
     if: github.repository == 'ultralytics/ultralytics' && (github.event_name == 'schedule' || github.event.inputs.raspberrypi == 'true')
-    timeout-minutes: 60
+    timeout-minutes: 120
     runs-on: raspberry-pi
     steps:
       - uses: actions/checkout@v4
@@ -253,7 +253,7 @@ jobs:
       - name: Install requirements
         run: |
           python -m pip install --upgrade pip wheel
-          pip install -e ".[export]" pytest mlflow pycocotools "ray[tune]"
+          pip install -e ".[export]" pytest
       - name: Check environment
         run: |
           yolo checks

tests/test_exports.py

@@ -23,22 +23,22 @@ from tests import MODEL, SOURCE
 
 def test_export_torchscript():
     """Test YOLO exports to TorchScript format."""
-    f = YOLO(MODEL).export(format="torchscript", optimize=False, imgsz=32)
+    file = YOLO(MODEL).export(format="torchscript", optimize=False, imgsz=32)
-    YOLO(f)(SOURCE, imgsz=32)  # exported model inference
+    YOLO(file)(SOURCE, imgsz=32)  # exported model inference
 
 
 def test_export_onnx():
     """Test YOLO exports to ONNX format."""
-    f = YOLO(MODEL).export(format="onnx", dynamic=True, imgsz=32)
+    file = YOLO(MODEL).export(format="onnx", dynamic=True, imgsz=32)
-    YOLO(f)(SOURCE, imgsz=32)  # exported model inference
+    YOLO(file)(SOURCE, imgsz=32)  # exported model inference
 
 
 @pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="OpenVINO not supported in Python 3.12")
 @pytest.mark.skipif(not TORCH_1_13, reason="OpenVINO requires torch>=1.13")
 def test_export_openvino():
     """Test YOLO exports to OpenVINO format."""
-    f = YOLO(MODEL).export(format="openvino", imgsz=32)
+    file = YOLO(MODEL).export(format="openvino", imgsz=32)
-    YOLO(f)(SOURCE, imgsz=32)  # exported model inference
+    YOLO(file)(SOURCE, imgsz=32)  # exported model inference
 
 
 @pytest.mark.slow
@@ -118,7 +118,7 @@ def test_export_torchscript_matrix(task, dynamic, int8, half, batch):
     ],
 )
 def test_export_coreml_matrix(task, dynamic, int8, half, batch):
-    """Test YOLO exports to TorchScript format."""
+    """Test YOLO exports to CoreML format."""
     file = YOLO(TASK2MODEL[task]).export(
         format="coreml",
         imgsz=32,
@@ -138,8 +138,8 @@ def test_export_coreml_matrix(task, dynamic, int8, half, batch):
 def test_export_coreml():
     """Test YOLO exports to CoreML format."""
     if MACOS:
-        f = YOLO(MODEL).export(format="coreml", imgsz=32)
+        file = YOLO(MODEL).export(format="coreml", imgsz=32)
-        YOLO(f)(SOURCE, imgsz=32)  # model prediction only supported on macOS for nms=False models
+        YOLO(file)(SOURCE, imgsz=32)  # model prediction only supported on macOS for nms=False models
     else:
         YOLO(MODEL).export(format="coreml", nms=True, imgsz=32)
 
@@ -152,8 +152,8 @@ def test_export_tflite():
     Note TF suffers from install conflicts on Windows and macOS.
     """
     model = YOLO(MODEL)
-    f = model.export(format="tflite", imgsz=32)
+    file = model.export(format="tflite", imgsz=32)
-    YOLO(f)(SOURCE, imgsz=32)
+    YOLO(file)(SOURCE, imgsz=32)
 
 
 @pytest.mark.skipif(True, reason="Test disabled")
@@ -165,8 +165,8 @@ def test_export_pb():
     Note TF suffers from install conflicts on Windows and macOS.
     """
     model = YOLO(MODEL)
-    f = model.export(format="pb", imgsz=32)
+    file = model.export(format="pb", imgsz=32)
-    YOLO(f)(SOURCE, imgsz=32)
+    YOLO(file)(SOURCE, imgsz=32)
 
 
 @pytest.mark.skipif(True, reason="Test disabled as Paddle protobuf and ONNX protobuf requirementsk conflict.")
@@ -182,5 +182,5 @@ def test_export_paddle():
 @pytest.mark.slow
 def test_export_ncnn():
     """Test YOLO exports to NCNN format."""
-    f = YOLO(MODEL).export(format="ncnn", imgsz=32)
+    file = YOLO(MODEL).export(format="ncnn", imgsz=32)
-    YOLO(f)(SOURCE, imgsz=32)  # exported model inference
+    YOLO(file)(SOURCE, imgsz=32)  # exported model inference

ultralytics/engine/exporter.py

@@ -83,6 +83,7 @@ from ultralytics.utils import (
     WINDOWS,
     __version__,
     callbacks,
+    checks,
     colorstr,
     get_default_args,
     yaml_save,
@@ -184,6 +185,7 @@ class Exporter:
         if sum(flags) != 1:
             raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
         jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn = flags  # export booleans
+        is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
 
         # Device
         if fmt == "engine" and self.args.device is None:
@@ -243,7 +245,7 @@ class Exporter:
                 m.dynamic = self.args.dynamic
                 m.export = True
                 m.format = self.args.format
-            elif isinstance(m, C2f) and not any((saved_model, pb, tflite, edgetpu, tfjs)):
+            elif isinstance(m, C2f) and not is_tf_format:
                 # EdgeTPU does not support FlexSplitV while split provides cleaner ONNX graph
                 m.forward = m.forward_split
 
@@ -303,7 +305,7 @@ class Exporter:
             f[3], _ = self.export_openvino()
         if coreml:  # CoreML
             f[4], _ = self.export_coreml()
-        if any((saved_model, pb, tflite, edgetpu, tfjs)):  # TensorFlow formats
+        if is_tf_format:  # TensorFlow formats
             self.args.int8 |= edgetpu
             f[5], keras_model = self.export_saved_model()
             if pb or tfjs:  # pb prerequisite to tfjs
@@ -777,11 +779,10 @@ class Exporter:
                     _ = self.cache.write_bytes(cache)
 
             # Load dataset w/ builder (for batching) and calibrate
-            dataset = self.get_int8_calibration_dataloader(prefix)
             config.int8_calibrator = EngineCalibrator(
-                dataset=dataset,
+                dataset=self.get_int8_calibration_dataloader(prefix),
                 batch=2 * self.args.batch,
-                cache=self.file.with_suffix(".cache"),
+                cache=str(self.file.with_suffix(".cache")),
             )
 
         elif half:
@@ -813,7 +814,7 @@ class Exporter:
         except ImportError:
             suffix = "-macos" if MACOS else "-aarch64" if ARM64 else "" if cuda else "-cpu"
             version = "" if ARM64 else "<=2.13.1"
-            check_requirements(f"tensorflow{suffix}{version}")
+            check_requirements((f"tensorflow{suffix}{version}", "keras"))
             import tensorflow as tf  # noqa
         if ARM64:
             check_requirements("cmake")  # 'cmake' is needed to build onnxsim on aarch64
@@ -855,24 +856,17 @@ class Exporter:
         f_onnx, _ = self.export_onnx()
 
         # Export to TF
-        tmp_file = f / "tmp_tflite_int8_calibration_images.npy"  # int8 calibration images file
         np_data = None
         if self.args.int8:
+            tmp_file = f / "tmp_tflite_int8_calibration_images.npy"  # int8 calibration images file
             verbosity = "info"
             if self.args.data:
-                # Generate calibration data for integer quantization
-                dataloader = self.get_int8_calibration_dataloader(prefix)
-                images = []
-                for i, batch in enumerate(dataloader):
-                    if i >= 100:  # maximum number of calibration images
-                        break
-                    im = batch["img"].permute(1, 2, 0)[None]  # list to nparray, CHW to BHWC
-                    images.append(im)
                 f.mkdir()
+                images = [batch["img"].permute(0, 2, 3, 1) for batch in self.get_int8_calibration_dataloader(prefix)]
                 images = torch.cat(images, 0).float()
                 # mean = images.view(-1, 3).mean(0)  # imagenet mean [123.675, 116.28, 103.53]
                 # std = images.view(-1, 3).std(0)  # imagenet std [58.395, 57.12, 57.375]
-                np.save(str(tmp_file), images.numpy())  # BHWC
+                np.save(str(tmp_file), images.numpy().astype(np.float32))  # BHWC
                 np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
         else:
             verbosity = "error"

ultralytics/utils/checks.py

@@ -23,7 +23,6 @@ from ultralytics.utils import (
     ASSETS,
     AUTOINSTALL,
     IS_COLAB,
-    IS_DOCKER,
     IS_JUPYTER,
     IS_KAGGLE,
     IS_PIP_PACKAGE,
@@ -322,17 +321,18 @@ def check_font(font="Arial.ttf"):
         return file
 
 
-def check_python(minimum: str = "3.8.0") -> bool:
+def check_python(minimum: str = "3.8.0", hard: bool = True) -> bool:
     """
     Check current python version against the required minimum version.
 
     Args:
         minimum (str): Required minimum version of python.
+        hard (bool, optional): If True, raise an AssertionError if the requirement is not met.
 
     Returns:
         (bool): Whether the installed Python version meets the minimum constraints.
     """
-    return check_version(PYTHON_VERSION, minimum, name="Python ", hard=True)
+    return check_version(PYTHON_VERSION, minimum, name="Python", hard=hard)
 
 
 @TryExcept()
@@ -735,4 +735,5 @@ def cuda_is_available() -> bool:
 
 
 # Define constants
+IS_PYTHON_MINIMUM_3_10 = check_python("3.10", hard=False)
 IS_PYTHON_3_12 = PYTHON_VERSION.startswith("3.12")