added dependencies

mct-2.1.1
Francesco Mattioli 2 months ago
parent bc49bcda0a
commit 24f2cfe61a
  1. 2
      ultralytics/engine/exporter.py
  2. 1
      ultralytics/nn/autobackend.py

@ -1041,7 +1041,7 @@ class Exporter:
@try_export @try_export
def export_mct(self, prefix=colorstr("Sony MCT:")): def export_mct(self, prefix=colorstr("Sony MCT:")):
check_requirements("mct-nightly") check_requirements(["mct-nightly","sony-custom-layers[torch]"])
import model_compression_toolkit as mct import model_compression_toolkit as mct
import onnx import onnx
from model_compression_toolkit.core import BitWidthConfig from model_compression_toolkit.core import BitWidthConfig

@ -178,6 +178,7 @@ class AutoBackend(nn.Module):
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if cuda else ["CPUExecutionProvider"] providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if cuda else ["CPUExecutionProvider"]
if mct: if mct:
check_requirements(["mct-nightly","sony-custom-layers[torch]"])
LOGGER.info(f"Loading {w} for ONNX MCT quantization inference...") LOGGER.info(f"Loading {w} for ONNX MCT quantization inference...")
import mct_quantizers as mctq import mct_quantizers as mctq
from sony_custom_layers.pytorch.object_detection import nms_ort # noqa from sony_custom_layers.pytorch.object_detection import nms_ort # noqa

Loading…
Cancel
Save