"""Callback that compresses a trained model by first exporting to .onnx format, and then converting to OpenVINO IR.""" # Copyright (C) 2020 Intel Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions # and limitations under the License. import logging import os from typing import Tuple from pytorch_lightning import Callback from anomalib.deploy import export_convert from anomalib.models.components import AnomalyModule logger = logging.getLogger(__name__) class OpenVINOCallback(Callback): """Callback to compresses a trained model. Model is first exported to ``.onnx`` format, and then converted to OpenVINO IR. Args: input_size (Tuple[int, int]): Tuple of image height, width dirpath (str): Path for model output filename (str): Name of output model """ def __init__(self, input_size: Tuple[int, int], dirpath: str, filename: str): self.input_size = input_size self.dirpath = dirpath self.filename = filename def on_train_end(self, trainer, pl_module: AnomalyModule) -> None: # pylint: disable=W0613 """Call when the train ends. Converts the model to ``onnx`` format and then calls OpenVINO's model optimizer to get the ``.xml`` and ``.bin`` IR files. """ logger.info("Exporting the model to OpenVINO") os.makedirs(self.dirpath, exist_ok=True) onnx_path = os.path.join(self.dirpath, self.filename + ".onnx") export_convert( model=pl_module, input_size=self.input_size, onnx_path=onnx_path, export_path=self.dirpath, )