def export(self, export_type: ExportType, output_model: ModelEntity):
assert export_type == ExportType.OPENVINO
output_model.model_format = ModelFormat.OPENVINO
output_model.optimization_type = self._optimization_type
with tempfile.TemporaryDirectory() as tempdir:
optimized_model_dir = os.path.join(tempdir, "dor")
logger.info(
f'Optimized model will be temporarily saved to "{optimized_model_dir}"'
)
os.makedirs(optimized_model_dir, exist_ok=True)
try:
onnx_model_path = os.path.join(optimized_model_dir,
'model.onnx')
mix_precision_status = self._model.mix_precision
self._model.mix_precision = False
export_onnx(self._model.eval(), self._cfg, onnx_model_path)
self._model.mix_precision = mix_precision_status
export_ir(onnx_model_path,
self._cfg.data.norm_mean,
self._cfg.data.norm_std,
optimized_model_dir=optimized_model_dir)
bin_file = [
f for f in os.listdir(optimized_model_dir)
if f.endswith('.bin')
][0]
xml_file = [
f for f in os.listdir(optimized_model_dir)
if f.endswith('.xml')
][0]
with open(os.path.join(optimized_model_dir, bin_file),
"rb") as f:
output_model.set_data("openvino.bin", f.read())
with open(os.path.join(optimized_model_dir, xml_file),
"rb") as f:
output_model.set_data("openvino.xml", f.read())
output_model.precision = self._precision
output_model.optimization_methods = self._optimization_methods
except Exception as ex:
raise RuntimeError('Optimization was unsuccessful.') from ex
output_model.set_data(
"label_schema.json",
label_schema_to_bytes(self._task_environment.label_schema))
logger.info('Exporting completed.')
def save_model(self, output_model: ModelEntity):
"""
Save the model after training is completed.
"""
config = self.get_config()
model_info = {
"model": self.model.state_dict(),
"config": config,
"VERSION": 1,
}
buffer = io.BytesIO()
torch.save(model_info, buffer)
output_model.set_data("weights.pth", buffer.getvalue())
output_model.set_data(
"label_schema.json",
label_schema_to_bytes(self.task_environment.label_schema))
# store computed threshold
output_model.set_data(
"threshold", bytes(struct.pack("f", self.model.threshold.item())))
f1_score = self.model.image_metrics.OptimalF1.compute().item()
output_model.performance = Performance(
score=ScoreMetric(name="F1 Score", value=f1_score))
output_model.precision = [ModelPrecision.FP32]
def optimize(self, optimization_type: OptimizationType,
dataset: DatasetEntity, output_model: ModelEntity,
optimization_parameters: Optional[OptimizationParameters]):
if optimization_type is not OptimizationType.POT:
raise ValueError(
"POT is the only supported optimization type for OpenVino models"
)
data_loader = OTEOpenVinoDataLoader(dataset, self.inferencer)
with tempfile.TemporaryDirectory() as tempdir:
xml_path = os.path.join(tempdir, "model.xml")
bin_path = os.path.join(tempdir, "model.bin")
with open(xml_path, "wb") as f:
f.write(self.model.get_data("openvino.xml"))
with open(bin_path, "wb") as f:
f.write(self.model.get_data("openvino.bin"))
model_config = ADDict({
'model_name': 'openvino_model',
'model': xml_path,
'weights': bin_path
})
model = load_model(model_config)
if get_nodes_by_type(model, ["FakeQuantize"]):
raise RuntimeError("Model is already optimized by POT")
engine_config = ADDict({'device': 'CPU'})
stat_subset_size = self.hparams.pot_parameters.stat_subset_size
preset = self.hparams.pot_parameters.preset.name.lower()
algorithms = [{
'name': 'DefaultQuantization',
'params': {
'target_device': 'ANY',
'preset': preset,
'stat_subset_size': min(stat_subset_size, len(data_loader)),
'shuffle_data': True
}
}]
engine = IEEngine(config=engine_config,
data_loader=data_loader,
metric=None)
pipeline = create_pipeline(algorithms, engine)
compressed_model = pipeline.run(model)
compress_model_weights(compressed_model)
with tempfile.TemporaryDirectory() as tempdir:
save_model(compressed_model, tempdir, model_name="model")
with open(os.path.join(tempdir, "model.xml"), "rb") as f:
output_model.set_data("openvino.xml", f.read())
with open(os.path.join(tempdir, "model.bin"), "rb") as f:
output_model.set_data("openvino.bin", f.read())
output_model.set_data(
"label_schema.json",
label_schema_to_bytes(self.task_environment.label_schema))
# set model attributes for quantized model
output_model.model_format = ModelFormat.OPENVINO
output_model.optimization_type = ModelOptimizationType.POT
output_model.optimization_methods = [OptimizationMethod.QUANTIZATION]
output_model.precision = [ModelPrecision.INT8]
self.model = output_model
self.inferencer = self.load_inferencer()
def optimize(
self,
optimization_type: OptimizationType,
dataset: DatasetEntity,
output_model: ModelEntity,
optimization_parameters: Optional[OptimizationParameters],
):
"""Optimize the model.
Args:
optimization_type (OptimizationType): Type of optimization [POT or NNCF]
dataset (DatasetEntity): Input Dataset.
output_model (ModelEntity): Output model.
optimization_parameters (Optional[OptimizationParameters]): Optimization parameters.
Raises:
ValueError: When the optimization type is not POT, which is the only support type at the moment.
"""
if optimization_type is not OptimizationType.POT:
raise ValueError(
"POT is the only supported optimization type for OpenVINO models"
)
data_loader = OTEOpenVINOAnomalyDataloader(config=self.config,
dataset=dataset,
inferencer=self.inferencer)
with tempfile.TemporaryDirectory() as tempdir:
xml_path = os.path.join(tempdir, "model.xml")
bin_path = os.path.join(tempdir, "model.bin")
self.__save_weights(
xml_path, self.task_environment.model.get_data("openvino.xml"))
self.__save_weights(
bin_path, self.task_environment.model.get_data("openvino.bin"))
model_config = {
"model_name": "openvino_model",
"model": xml_path,
"weights": bin_path,
}
model = load_model(model_config)
if get_nodes_by_type(model, ["FakeQuantize"]):
logger.warning("Model is already optimized by POT")
return
engine = IEEngine(config=ADDict({"device": "CPU"}),
data_loader=data_loader,
metric=None)
pipeline = create_pipeline(
algo_config=self._get_optimization_algorithms_configs(),
engine=engine)
compressed_model = pipeline.run(model)
compress_model_weights(compressed_model)
with tempfile.TemporaryDirectory() as tempdir:
save_model(compressed_model, tempdir, model_name="model")
self.__load_weights(path=os.path.join(tempdir, "model.xml"),
output_model=output_model,
key="openvino.xml")
self.__load_weights(path=os.path.join(tempdir, "model.bin"),
output_model=output_model,
key="openvino.bin")
output_model.set_data(
"label_schema.json",
label_schema_to_bytes(self.task_environment.label_schema))
output_model.set_data(
"threshold", self.task_environment.model.get_data("threshold"))
output_model.model_status = ModelStatus.SUCCESS
output_model.model_format = ModelFormat.OPENVINO
output_model.optimization_type = ModelOptimizationType.POT
output_model.optimization_methods = [OptimizationMethod.QUANTIZATION]
output_model.precision = [ModelPrecision.INT8]
self.task_environment.model = output_model
self.inferencer = self.load_inferencer()
def optimize(
self,
optimization_type: OptimizationType,
dataset: DatasetEntity,
output_model: ModelEntity,
optimization_parameters: Optional[OptimizationParameters],
):
""" Optimize a model on a dataset """
if optimization_type is not OptimizationType.NNCF:
raise RuntimeError('NNCF is the only supported optimization')
if self._compression_ctrl:
raise RuntimeError('The model is already optimized. NNCF requires the original model for optimization.')
if self._cfg.train.ema.enable:
raise RuntimeError('EMA model could not be used together with NNCF compression')
if self._cfg.lr_finder.enable:
raise RuntimeError('LR finder could not be used together with NNCF compression')
aux_pretrained_dicts = self._load_aux_models_data(self._task_environment.model)
num_aux_models = len(self._cfg.mutual_learning.aux_configs)
num_aux_pretrained_dicts = len(aux_pretrained_dicts)
if num_aux_models != num_aux_pretrained_dicts:
raise RuntimeError('The pretrained weights are not provided for all aux models.')
if optimization_parameters is not None:
update_progress_callback = optimization_parameters.update_progress
else:
update_progress_callback = default_progress_callback
time_monitor = TrainingProgressCallback(update_progress_callback, num_epoch=self._cfg.train.max_epoch,
num_train_steps=math.ceil(len(dataset.get_subset(Subset.TRAINING)) /
self._cfg.train.batch_size),
num_val_steps=0, num_test_steps=0)
self.metrics_monitor = DefaultMetricsMonitor()
self.stop_callback.reset()
set_random_seed(self._cfg.train.seed)
train_subset = dataset.get_subset(Subset.TRAINING)
val_subset = dataset.get_subset(Subset.VALIDATION)
self._cfg.custom_datasets.roots = [OTEClassificationDataset(train_subset, self._labels, self._multilabel,
keep_empty_label=self._empty_label in self._labels),
OTEClassificationDataset(val_subset, self._labels, self._multilabel,
keep_empty_label=self._empty_label in self._labels)]
datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(self._cfg))
self._compression_ctrl, self._model, self._nncf_metainfo = \
wrap_nncf_model(self._model, self._cfg, datamanager_for_init=datamanager)
self._cfg.train.lr = calculate_lr_for_nncf_training(self._cfg, self._initial_lr, False)
train_model = self._model
if self._cfg.use_gpu:
main_device_ids = list(range(self.num_devices))
extra_device_ids = [main_device_ids for _ in range(num_aux_models)]
train_model = DataParallel(train_model, device_ids=main_device_ids,
output_device=0).cuda(main_device_ids[0])
else:
extra_device_ids = [None for _ in range(num_aux_models)]
optimizer = torchreid.optim.build_optimizer(train_model, **optimizer_kwargs(self._cfg))
scheduler = torchreid.optim.build_lr_scheduler(optimizer, num_iter=datamanager.num_iter,
**lr_scheduler_kwargs(self._cfg))
logger.info('Start training')
run_training(self._cfg, datamanager, train_model, optimizer,
scheduler, extra_device_ids, self._cfg.train.lr,
should_freeze_aux_models=True,
aux_pretrained_dicts=aux_pretrained_dicts,
tb_writer=self.metrics_monitor,
perf_monitor=time_monitor,
stop_callback=self.stop_callback,
nncf_metainfo=self._nncf_metainfo,
compression_ctrl=self._compression_ctrl)
self.metrics_monitor.close()
if self.stop_callback.check_stop():
logger.info('Training cancelled.')
return
logger.info('Training completed')
self.save_model(output_model)
output_model.model_format = ModelFormat.BASE_FRAMEWORK
output_model.optimization_type = self._optimization_type
output_model.optimization_methods = self._optimization_methods
output_model.precision = self._precision