Chiebot-Mirror
/
ultralytics
mirror of https://github.com/ultralytics/ultralytics.git
9
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

Auto-format by https://ultralytics.com/actions

Browse Source
mct-2.1.1
UltralyticsAssistant 6 months ago
parent adaee64f8f
commit 9a5fcfdadc
2 changed files with 58 additions and 57 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 95
      ultralytics/engine/exporter.py
  2. 20
      ultralytics/nn/autobackend.py

95
ultralytics/engine/exporter.py
Unescape View File

@ -184,7 +184,9 @@ class Exporter:
flags = [x == fmt for x in fmts] flags = [x == fmt for x in fmts]
if sum(flags) != 1: if sum(flags) != 1:
raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}") raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn, mct = flags # export booleans jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn, mct = (
flags # export booleans
)
is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs)) is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
if mct: if mct:
LOGGER.warning("WARNING ⚠️ Sony MCT only supports int8 export, setting int8=True.") LOGGER.warning("WARNING ⚠️ Sony MCT only supports int8 export, setting int8=True.")
@ -1020,21 +1022,24 @@ class Exporter:
# j.write(subst) # j.write(subst)
yaml_save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml yaml_save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
return f, None return f, None
@try_export @try_export
def export_mct(self, prefix=colorstr("Sony MCT:")): def export_mct(self, prefix=colorstr("Sony MCT:")):
# pip install --upgrade -force-reinstall git+https://github.com/ambitious-octopus/model_optimization.git@get-output-fix # pip install --upgrade -force-reinstall git+https://github.com/ambitious-octopus/model_optimization.git@get-output-fix
import model_compression_toolkit as mct import model_compression_toolkit as mct
from torch import nn
# pip install sony-custom-layers[torch] # pip install sony-custom-layers[torch]
from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms
from torch import nn
class PostProcessWrapper(nn.Module): class PostProcessWrapper(nn.Module):
def __init__(self, def __init__(
model: nn.Module, self,
score_threshold: float = 0.001, model: nn.Module,
iou_threshold: float = 0.7, score_threshold: float = 0.001,
max_detections: int = 300): iou_threshold: float = 0.7,
max_detections: int = 300,
):
""" """
Wrapping PyTorch Module with multiclass_nms layer from sony_custom_layers. Wrapping PyTorch Module with multiclass_nms layer from sony_custom_layers.
@ -1056,49 +1061,59 @@ class Exporter:
boxes = outputs[0] boxes = outputs[0]
scores = outputs[1] scores = outputs[1]
nms = multiclass_nms(boxes=boxes, scores=scores, score_threshold=self.score_threshold, nms = multiclass_nms(
iou_threshold=self.iou_threshold, max_detections=self.max_detections) boxes=boxes,
scores=scores,
score_threshold=self.score_threshold,
iou_threshold=self.iou_threshold,
max_detections=self.max_detections,
)
return nms return nms
def representative_dataset_gen(dataloader=self.get_int8_calibration_dataloader(prefix)): def representative_dataset_gen(dataloader=self.get_int8_calibration_dataloader(prefix)):
for batch in dataloader: for batch in dataloader:
img = batch["img"] img = batch["img"]
img = img / 255.0 img = img / 255.0
yield [img] yield [img]
tpc = mct.get_target_platform_capabilities(fw_name="pytorch", tpc = mct.get_target_platform_capabilities(
target_platform_name='imx500', fw_name="pytorch", target_platform_name="imx500", target_platform_version="v1"
target_platform_version='v1') )
mp_config = mct.core.MixedPrecisionQuantizationConfig(num_of_images=5, mp_config = mct.core.MixedPrecisionQuantizationConfig(num_of_images=5, use_hessian_based_scores=False)
use_hessian_based_scores=False) config = mct.core.CoreConfig(
config = mct.core.CoreConfig(mixed_precision_config=mp_config, mixed_precision_config=mp_config,
quantization_config=mct.core.QuantizationConfig(shift_negative_activation_correction=True)) quantization_config=mct.core.QuantizationConfig(shift_negative_activation_correction=True),
)
resource_utilization_data = mct.core.pytorch_resource_utilization_data(in_model=self.model,
representative_data_gen= resource_utilization_data = mct.core.pytorch_resource_utilization_data(
representative_dataset_gen, in_model=self.model,
core_config=config, representative_data_gen=representative_dataset_gen,
target_platform_capabilities=tpc) core_config=config,
target_platform_capabilities=tpc,
resource_utilization = mct.core.ResourceUtilization(weights_memory=resource_utilization_data.weights_memory * 0.75) )
quant_model, _ = mct.ptq.pytorch_post_training_quantization(in_module=self.model, resource_utilization = mct.core.ResourceUtilization(
representative_data_gen= weights_memory=resource_utilization_data.weights_memory * 0.75
representative_dataset_gen, )
target_resource_utilization=resource_utilization,
core_config=config, quant_model, _ = mct.ptq.pytorch_post_training_quantization(
target_platform_capabilities=tpc) in_module=self.model,
representative_data_gen=representative_dataset_gen,
target_resource_utilization=resource_utilization,
core_config=config,
target_platform_capabilities=tpc,
)
# Get working device # Get working device
device = mct.core.pytorch.pytorch_device_config.get_working_device() device = mct.core.pytorch.pytorch_device_config.get_working_device()
quant_model_pp = PostProcessWrapper(model=quant_model).to(device=device) quant_model_pp = PostProcessWrapper(model=quant_model).to(device=device)
f = str(self.file).replace(self.file.suffix, "_mct_model.onnx") f = str(self.file).replace(self.file.suffix, "_mct_model.onnx")
mct.exporter.pytorch_export_model(model=quant_model_pp, mct.exporter.pytorch_export_model(
save_model_path=f, model=quant_model_pp, save_model_path=f, repr_dataset=representative_dataset_gen
repr_dataset=representative_dataset_gen) )
return f, None return f, None
def _add_tflite_metadata(self, file): def _add_tflite_metadata(self, file):
"""Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata.""" """Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata."""
import flatbuffers import flatbuffers

20
ultralytics/nn/autobackend.py
Unescape View File

@ -106,23 +106,9 @@ class AutoBackend(nn.Module):
super().__init__() super().__init__()
w = str(weights[0] if isinstance(weights, list) else weights) w = str(weights[0] if isinstance(weights, list) else weights)
nn_module = isinstance(weights, torch.nn.Module) nn_module = isinstance(weights, torch.nn.Module)
( (pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn, triton, mct) = (
pt, self._model_type(w)
jit, )
onnx,
xml,
engine,
coreml,
saved_model,
pb,
tflite,
edgetpu,
tfjs,
paddle,
ncnn,
triton,
mct
) = self._model_type(w)
fp16 &= pt or jit or onnx or xml or engine or nn_module or triton # FP16 fp16 &= pt or jit or onnx or xml or engine or nn_module or triton # FP16
nhwc = coreml or saved_model or pb or tflite or edgetpu # BHWC formats (vs torch BCWH) nhwc = coreml or saved_model or pb or tflite or edgetpu # BHWC formats (vs torch BCWH)
stride = 32 # default stride stride = 32 # default stride

Write Preview
Loading…
Cancel
Save