77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
# Copyright (C) 2020-2022 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import json
|
|
import numpy as np
|
|
|
|
import pytest
|
|
from addict import Dict
|
|
|
|
from openvino.tools.pot.graph import load_model
|
|
from openvino.tools.pot.graph.model_utils import get_node_by_name
|
|
from openvino.tools.pot.graph import node_utils as nu
|
|
from openvino.tools.pot.algorithms.quantization import fake_quantize as fqut
|
|
from openvino.tools.pot.pipeline.initializer import create_pipeline
|
|
from openvino.tools.pot.engines.ac_engine import ACEngine
|
|
from .utils.config import get_engine_config, merge_configs
|
|
from .utils.path import TEST_ROOT
|
|
|
|
TEST_MODELS = [
|
|
('mobilenet-v2-pytorch', 'pytorch', 'MinMaxQuantization', 'performance'),
|
|
('resnet-50-tf', 'tf', 'DefaultQuantization', 'performance'),
|
|
('octave-resnet-26-0.25', 'mxnet', 'DefaultQuantization', 'accuracy'),
|
|
]
|
|
|
|
REFERENCES_PATH = TEST_ROOT / 'data/reference_unify'
|
|
|
|
|
|
@pytest.fixture(scope='module', params=TEST_MODELS,
|
|
ids=["{}_{}_{}_{}".format(*m) for m in TEST_MODELS])
|
|
def _params(request):
|
|
return request.param
|
|
|
|
|
|
def test_unify_scales(_params, tmp_path, models):
|
|
model_name, model_framework, algorithm, preset = _params
|
|
|
|
algorithm_config = Dict({
|
|
'algorithms': [{
|
|
'name': algorithm,
|
|
'params': {
|
|
'target_device': 'VPU',
|
|
'preset': preset,
|
|
'stat_subset_size': 2
|
|
}
|
|
}]
|
|
})
|
|
|
|
def _test_unify_scales(model_, to_unify_):
|
|
for _, fqs in to_unify_:
|
|
ranges = []
|
|
for fq in fqs:
|
|
fq = get_node_by_name(model_, fq)
|
|
fq_inputs = nu.get_node_inputs(fq)[1:]
|
|
ranges.append(tuple(fqut.get_node_value(fq_input) for fq_input in fq_inputs))
|
|
assert all([np.array_equal(r, ranges[0][i]) for i, r in enumerate(ranges[-1])])
|
|
|
|
model = models.get(model_name, model_framework, tmp_path)
|
|
|
|
engine_config = get_engine_config(model_name)
|
|
config = merge_configs(model.model_params, engine_config, algorithm_config)
|
|
|
|
model = load_model(config.model)
|
|
pipeline = create_pipeline(config.compression.algorithms, ACEngine(config.engine))
|
|
compressed_model = pipeline.run(model)
|
|
|
|
to_unify = fqut.find_fqs_to_unify(compressed_model, config.compression.algorithms[0]['params'])
|
|
_test_unify_scales(compressed_model, to_unify)
|
|
|
|
ref_path = REFERENCES_PATH.joinpath(model_name + '_to_unify.json')
|
|
if ref_path.exists():
|
|
with open(ref_path.as_posix(), 'r') as f:
|
|
to_unify_ref = json.load(f)
|
|
assert to_unify == to_unify_ref
|
|
else:
|
|
with open(ref_path.as_posix(), 'w+') as f:
|
|
json.dump(to_unify, f, indent=4)
|