From 11623eeb002730553c655a9d873f495d233290e0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 24 May 2024 01:00:17 +0200 Subject: [PATCH] Fix TFLite INT8 quant bug (#13082) --- .github/workflows/ci.yaml | 6 +++--- tests/test_exports.py | 30 +++++++++++++++--------------- ultralytics/engine/exporter.py | 26 ++++++++++---------------- ultralytics/utils/checks.py | 7 ++++--- 4 files changed, 32 insertions(+), 37 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index f96a967d..373303f7 100644 --- a/.github/workflows/ci.yaml +++ b/.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 diff --git a/tests/test_exports.py b/tests/test_exports.py index 788861e6..2a500a71 100644 --- a/tests/test_exports.py +++ b/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) - YOLO(f)(SOURCE, imgsz=32) # exported model inference + file = YOLO(MODEL).export(format="torchscript", optimize=False, imgsz=32) + 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) - YOLO(f)(SOURCE, imgsz=32) # exported model inference + file = YOLO(MODEL).export(format="onnx", dynamic=True, imgsz=32) + 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) - YOLO(f)(SOURCE, imgsz=32) # exported model inference + file = YOLO(MODEL).export(format="openvino", imgsz=32) + 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) - YOLO(f)(SOURCE, imgsz=32) # model prediction only supported on macOS for nms=False models + file = YOLO(MODEL).export(format="coreml", imgsz=32) + 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) - YOLO(f)(SOURCE, imgsz=32) + file = model.export(format="tflite", 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) - YOLO(f)(SOURCE, imgsz=32) + file = model.export(format="pb", 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) - YOLO(f)(SOURCE, imgsz=32) # exported model inference + file = YOLO(MODEL).export(format="ncnn", imgsz=32) + YOLO(file)(SOURCE, imgsz=32) # exported model inference diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index cdb635e3..f050d5ea 100644 --- a/ultralytics/engine/exporter.py +++ b/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" diff --git a/ultralytics/utils/checks.py b/ultralytics/utils/checks.py index e5290e18..6f2807d3 100644 --- a/ultralytics/utils/checks.py +++ b/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")