From 45ba99973d8a578c88533a9da91887dddb9de3f6 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 31 Aug 2023 15:19:44 +0200 Subject: [PATCH] Add NNCF support in OpenVINO export (#4671) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 下北泽miHoMo红茶坊 <39751846+kisaragychihaya@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/docker.yaml | 37 ++++++++++-------- docker/Dockerfile-cpu | 4 +- ultralytics/data/augment.py | 7 +--- ultralytics/engine/exporter.py | 69 +++++++++++++++++++++++----------- 4 files changed, 72 insertions(+), 45 deletions(-) diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml index faf0bcc9..7d0c8e4d 100644 --- a/.github/workflows/docker.yaml +++ b/.github/workflows/docker.yaml @@ -8,18 +8,25 @@ on: branches: [main] workflow_dispatch: inputs: - dockerfile: - type: choice - description: Select Dockerfile - options: - - Dockerfile-arm64 - - Dockerfile-jetson - - Dockerfile-python - - Dockerfile-cpu - - Dockerfile + Dockerfile: + type: boolean + description: Use Dockerfile + default: true + Dockerfile-cpu: + type: boolean + description: Use Dockerfile-cpu + Dockerfile-arm64: + type: boolean + description: Use Dockerfile-arm64 + Dockerfile-jetson: + type: boolean + description: Use Dockerfile-jetson + Dockerfile-python: + type: boolean + description: Use Dockerfile-python push: type: boolean - description: Push image to Docker Hub + description: Push images to Docker Hub default: true jobs: @@ -93,28 +100,28 @@ jobs: VERSION_TAG: ${{ steps.get_version.outputs.version_tag }} - name: Build Image - if: github.event_name == 'push' || github.event.inputs.dockerfile == matrix.dockerfile + if: github.event_name == 'push' || github.event.inputs[matrix.dockerfile] == 'true' run: | docker build --platform ${{ matrix.platforms }} -f docker/${{ matrix.dockerfile }} \ -t ultralytics/ultralytics:${{ matrix.tags }} \ -t ultralytics/ultralytics:${{ steps.get_version.outputs.version_tag }} . - name: Run Tests - if: (github.event_name == 'push' || github.event.inputs.dockerfile == matrix.dockerfile) && matrix.platforms == 'linux/amd64' # arm64 images not supported on GitHub CI runners + if: (github.event_name == 'push' || github.event.inputs[matrix.dockerfile] == 'true') && matrix.platforms == 'linux/amd64' # arm64 images not supported on GitHub CI runners run: docker run ultralytics/ultralytics:${{ matrix.tags }} /bin/bash -c "pip install pytest && pytest tests" - name: Run Benchmarks # WARNING: Dockerfile (GPU) error on TF.js export 'module 'numpy' has no attribute 'object'. - if: (github.event_name == 'push' || github.event.inputs.dockerfile == matrix.dockerfile) && matrix.platforms == 'linux/amd64' && matrix.dockerfile != 'Dockerfile' # arm64 images not supported on GitHub CI runners + if: (github.event_name == 'push' || github.event.inputs[matrix.dockerfile] == 'true') && matrix.platforms == 'linux/amd64' && matrix.dockerfile != 'Dockerfile' # arm64 images not supported on GitHub CI runners run: docker run ultralytics/ultralytics:${{ matrix.tags }} yolo benchmark model=yolov8n.pt imgsz=160 verbose=0.26 - name: Push Docker Image with Ultralytics version tag - if: (github.event_name == 'push' || (github.event.inputs.dockerfile == matrix.dockerfile && github.event.inputs.push == 'true')) && steps.check_tag.outputs.exists == 'false' + if: (github.event_name == 'push' || (github.event.inputs[matrix.dockerfile] == 'true' && github.event.inputs.push == 'true')) && steps.check_tag.outputs.exists == 'false' run: | docker push ultralytics/ultralytics:${{ steps.get_version.outputs.version_tag }} - name: Push Docker Image with latest tag - if: github.event_name == 'push' || (github.event.inputs.dockerfile == matrix.dockerfile && github.event.inputs.push == 'true') + if: github.event_name == 'push' || (github.event.inputs[matrix.dockerfile] == 'true' && github.event.inputs.push == 'true') run: | docker push ultralytics/ultralytics:${{ matrix.tags }} if [[ "${{ matrix.tags }}" == "latest" ]]; then diff --git a/docker/Dockerfile-cpu b/docker/Dockerfile-cpu index 3e635963..7879affd 100644 --- a/docker/Dockerfile-cpu +++ b/docker/Dockerfile-cpu @@ -18,8 +18,8 @@ RUN apt update \ WORKDIR /usr/src/ultralytics # Copy contents -# COPY . /usr/src/app (issues as not a .git directory) -RUN git clone https://github.com/ultralytics/ultralytics /usr/src/ultralytics +COPY . /usr/src/ultralytics +# RUN git clone https://github.com/ultralytics/ultralytics /usr/src/ultralytics ADD https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt /usr/src/ultralytics/ # Remove python3.11/EXTERNALLY-MANAGED or use 'pip install --break-system-packages' avoid 'externally-managed-environment' Ubuntu nightly error diff --git a/ultralytics/data/augment.py b/ultralytics/data/augment.py index d24f7982..f14b82bc 100644 --- a/ultralytics/data/augment.py +++ b/ultralytics/data/augment.py @@ -65,12 +65,7 @@ class Compose: def __repr__(self): """Return string representation of object.""" - format_string = f'{self.__class__.__name__}(' - for t in self.transforms: - format_string += '\n' - format_string += f' {t}' - format_string += '\n)' - return format_string + return f"{self.__class__.__name__}({', '.join([f'{t}' for t in self.transforms])})" class BaseMixTransform: diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index c81099e8..defa2d06 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -58,9 +58,12 @@ from copy import deepcopy from datetime import datetime from pathlib import Path +import numpy as np import torch from ultralytics.cfg import get_cfg +from ultralytics.data.dataset import YOLODataset +from ultralytics.data.utils import check_det_dataset from ultralytics.nn.autobackend import check_class_names from ultralytics.nn.modules import C2f, Detect, RTDETRDecoder from ultralytics.nn.tasks import DetectionModel, SegmentationModel @@ -275,10 +278,11 @@ class Exporter: f"work. Use export 'imgsz={max(self.imgsz)}' if val is required." imgsz = self.imgsz[0] if square else str(self.imgsz)[1:-1].replace(' ', '') predict_data = f'data={data}' if model.task == 'segment' and format == 'pb' else '' + q = 'int8' if self.args.int8 else 'half' if self.args.half else '' # quantization LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)' f"\nResults saved to {colorstr('bold', file.parent.resolve())}" - f'\nPredict: yolo predict task={model.task} model={f} imgsz={imgsz} {predict_data}' - f'\nValidate: yolo val task={model.task} model={f} imgsz={imgsz} data={data} {s}' + f'\nPredict: yolo predict task={model.task} model={f} imgsz={imgsz} {q} {predict_data}' + f'\nValidate: yolo val task={model.task} model={f} imgsz={imgsz} data={data} {q} {s}' f'\nVisualize: https://netron.app') self.run_callbacks('on_export_end') @@ -367,27 +371,54 @@ class Exporter: LOGGER.info(f'\n{prefix} starting export with openvino {ov.__version__}...') f = str(self.file).replace(self.file.suffix, f'_openvino_model{os.sep}') + fq = str(self.file).replace(self.file.suffix, f'_int8_openvino_model{os.sep}') f_onnx = self.file.with_suffix('.onnx') f_ov = str(Path(f) / self.file.with_suffix('.xml').name) + fq_ov = str(Path(fq) / self.file.with_suffix('.xml').name) + + def serialize(ov_model, file): + """Set RT info, serialize and save metadata YAML.""" + ov_model.set_rt_info('YOLOv8', ['model_info', 'model_type']) + ov_model.set_rt_info(True, ['model_info', 'reverse_input_channels']) + ov_model.set_rt_info(114, ['model_info', 'pad_value']) + ov_model.set_rt_info([255.0], ['model_info', 'scale_values']) + ov_model.set_rt_info(self.args.iou, ['model_info', 'iou_threshold']) + ov_model.set_rt_info([v.replace(' ', '_') for v in self.model.names.values()], ['model_info', 'labels']) + if self.model.task != 'classify': + ov_model.set_rt_info('fit_to_window_letterbox', ['model_info', 'resize_type']) + + ov.serialize(ov_model, file) # save + yaml_save(Path(file).parent / 'metadata.yaml', self.metadata) # add metadata.yaml ov_model = mo.convert_model(f_onnx, model_name=self.pretty_name, framework='onnx', compress_to_fp16=self.args.half) # export - # Set RT info - ov_model.set_rt_info('YOLOv8', ['model_info', 'model_type']) - ov_model.set_rt_info(True, ['model_info', 'reverse_input_channels']) - ov_model.set_rt_info(114, ['model_info', 'pad_value']) - ov_model.set_rt_info([255.0], ['model_info', 'scale_values']) - ov_model.set_rt_info(self.args.iou, ['model_info', 'iou_threshold']) - ov_model.set_rt_info([v.replace(' ', '_') for k, v in sorted(self.model.names.items())], - ['model_info', 'labels']) - if self.model.task != 'classify': - ov_model.set_rt_info('fit_to_window_letterbox', ['model_info', 'resize_type']) + if self.args.int8: + assert self.args.data, "INT8 export requires a data argument for calibration, i.e. 'data=coco8.yaml'" + check_requirements('nncf>=2.5.0') + import nncf - ov.serialize(ov_model, f_ov) # save - yaml_save(Path(f) / 'metadata.yaml', self.metadata) # add metadata.yaml + def transform_fn(data_item): + """Quantization transform function.""" + im = data_item['img'].numpy().astype(np.float32) / 255.0 # uint8 to fp16/32 and 0 - 255 to 0.0 - 1.0 + return np.expand_dims(im, 0) if im.ndim == 3 else im + + # 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) + quantization_dataset = nncf.Dataset(dataset, transform_fn) + ignored_scope = nncf.IgnoredScope(types=['Multiply', 'Subtract', 'Sigmoid']) # ignore operation + quantized_ov_model = nncf.quantize(ov_model, + quantization_dataset, + preset=nncf.QuantizationPreset.MIXED, + ignored_scope=ignored_scope) + serialize(quantized_ov_model, fq_ov) + return fq, None + + serialize(ov_model, f_ov) return f, None @try_export @@ -633,19 +664,13 @@ class Exporter: if self.args.int8: verbosity = '--verbosity info' if self.args.data: - import numpy as np - - from ultralytics.data.dataset import YOLODataset - from ultralytics.data.utils import check_det_dataset - # 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) images = [] - n_images = 100 # maximum number of images - for n, batch in enumerate(dataset): - if n >= n_images: + for i, batch in enumerate(dataset): + 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)