mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 05:24:22 +08:00
Add NNCF support in OpenVINO export (#4671)
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>
This commit is contained in:
parent
d2cf7acce0
commit
45ba99973d
37
.github/workflows/docker.yaml
vendored
37
.github/workflows/docker.yaml
vendored
@ -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
|
||||
|
@ -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
|
||||
|
@ -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:
|
||||
|
@ -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)
|
||||
|
Loading…
x
Reference in New Issue
Block a user