#########################################################
# Copyright (C) 2020 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Joey Chou
#########################################################
import os
import yaml
from typing import Dict, TypeVar, Tuple, Optional, List, Iterable, Any
from sima_utils.data.data_generator import DataGenerator
from afe.backends import Backend
from afe.core.configs import (
ModelConfigs, QuantizationConfigs, CompressionConfigs, OptimizationConfigs,
TransformerConfigs, api_calibration_configs, create_quantization_configs)
from afe.ir.defines import InputShape
from afe.ir.net import AwesomeNet
from afe.ir.serializer.api import save_awesomenet
[docs]
def parse_yaml(yaml_filepath: str) -> Dict:
with open(yaml_filepath, 'r') as stream:
config = yaml.safe_load(stream)
return config
[docs]
def dataclass_to_dict(_class):
from enum import Enum
if isinstance(_class, Enum):
return _class.value
try:
_class.__dict__
except:
return _class
_dict = {}
for k, v in _class.__dict__.items():
_dict[k] = dataclass_to_dict(v)
return _dict
[docs]
def load_configs_from_yaml(file_path: str) -> Tuple[ModelConfigs, OptimizationConfigs]:
with open(file_path) as f:
loaded = yaml.load(f, Loader=yaml.FullLoader)
model_config = ModelConfigs(**loaded["Model Configurations"])
opt_config_dict = loaded["Optimization Configurations"]
quant_config = QuantizationConfigs(**opt_config_dict["quantization_configs"])
comp_config = CompressionConfigs(**opt_config_dict["compression_configs"])
opt_config_dict["quantization_configs"] = quant_config
opt_config_dict["compression_configs"] = comp_config
opt_config = OptimizationConfigs(**opt_config_dict)
return model_config, opt_config
[docs]
def dump_configs_to_yaml(model_config: ModelConfigs,
opt_config: OptimizationConfigs) -> None:
"""
Dump the YAML file containing ModelConfigs and OptimizationConfigs to directory:
{model_config.output_dir}/{model_config.model_name}.yaml
"""
yaml_dict = {"Model Configurations": dataclass_to_dict(model_config),
"Optimization Configurations": dataclass_to_dict(opt_config)}
yaml_file_name = model_config.output_directory + "/" + model_config.name + ".yaml"
with open(yaml_file_name, 'w') as f:
yaml.dump(yaml_dict, f)
[docs]
def dump_yaml_npz(model_config: ModelConfigs,
net: Optional[AwesomeNet] = None,
name_prefix: str = "",
name_postfix: str = "") -> None:
"""
Dump the yaml and npz to directory:
{model_config.output_dir}/{name_prefix}{model_config.model_name}{name_postfix}.yaml
{model_config.output_dir}/{name_prefix}{model_config.model_name}{name_postfix}.npz
"""
filename = f"{name_prefix}{model_config.name}{name_postfix}"
save_awesomenet(net=net,
model_name=filename,
output_directory=model_config.output_directory)
class _DataGeneratorIterableProxy(Iterable):
"""
A wrapper around DataGenerator that fixes discrepancies in
DataGenerator's iterable interface.
Iteration returns the same values as the data generator's __getitem__ method.
:param g: Data generator to wrap
:param length_limit: If not None, the maximum number of items to use from the
data generator. Excess items are ignored.
"""
def __init__(self, g: DataGenerator, *, length_limit: Optional[int] = None):
self._generator = g
length = len(g)
if length_limit is not None:
length = min(length, length_limit)
self._length = length
def __iter__(self):
return (self._generator[i] for i in range(self._length))
def __len__(self):
return self._length
[docs]
def convert_data_generator_to_iterable(g: DataGenerator, *, length_limit: Optional[int] = None) -> Iterable[Any]:
"""
Convert a data generator to an iterable object.
Although DataGenerator has methods like an iterable object, it does not
implement the iterable interface properly.
:param g: Data generator
:param length_limit: If not None, the maximum number of items to use from the
data generator. Excess items are ignored.
:return: Iterable over the sequence g[0], g[1], .... The data generator
must not be modified while the iterable is being used.
"""
return _DataGeneratorIterableProxy(g, length_limit=length_limit)
[docs]
class LengthHintedIterable(Iterable[T]):
"""
Wrapper class, which wraps Iterables and their length hints. Used for length hint in our API.
It is intended to be instantiated with method length_hinted().
model = loaded_net.quantize(length_hinted(24, data_source), default_quantization)
"""
_length_hint: int
_iterable: Iterable[T]
def __init__(self, length_hint: int, iterable: Iterable[T]):
self._length_hint = length_hint
self._iterable = iterable
def __iter__(self):
return self._iterable.__iter__()
[docs]
def get_length(self):
return self._length_hint
[docs]
def length_hinted(length_hint: int, iterable: Iterable[T]) -> LengthHintedIterable[T]:
"""
Used to create Length hinted iterable, and use it on our API. Example of usage:
model = loaded_net.quantize(length_hinted(24, data_source), default_quantization)
:param length_hint: Specified number of examples
:param iterable: Input examples as Iterable.
"""
return LengthHintedIterable(length_hint, iterable)
[docs]
def save_files() -> bool:
return os.environ.get('SIMA_AFE_SAVED_FILES') == '1'
[docs]
def wrap_parameters_to_model_configs(name: str,
framework: str = "",
input_names: Optional[List[str]] = None,
input_shapes: Optional[List[InputShape]] = None,
input_dtypes: Optional[List[str]] = None,
layout: str = "",
model_path: str = "",
model_file_paths: Optional[List[str]] = None,
output_names: Optional[List[str]] = None,
output_directory: Optional[str] = None,
is_quantized: bool = False
) -> ModelConfigs:
"""
Given the list of model parameters, create ModelConfigs data structure.
:param name: str. Model name.
:param framework: str. Framework used in a model.
:param input_names: Optional[List[str]]. List of input names to a model, if any.
:param input_shapes: Optional[List[InputShape]]. List of input shapes, if any.
:param input_dtypes: Optional[List[str]]. List of input types to a model, if any.
:param layout: str. Data layout used in a model.
:param model_path: str. The file path from which the model is loaded.
:param model_file_paths: Optional[List[str]]. The file paths used for model loading. Used in cases where
multiple files are needed to load a model.
:param output_names: Optional[List[str]]. List of output names, if any.
:param output_directory: Optional[str]. Output directory path used to store generated files, if any.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: ModelConfigs.
"""
input_names_list: List[str] = input_names if input_names is not None else []
input_shapes_list: List[InputShape] = input_shapes if input_shapes is not None else []
input_dtypes_list: List[str] = input_dtypes if input_dtypes is not None else []
model_file_paths_list: List[str] = model_file_paths if model_file_paths is not None else []
return ModelConfigs(name, framework, input_names_list, input_shapes_list, input_dtypes_list, layout,
model_path, model_file_paths_list, is_quantized, output_names, output_directory)
[docs]
def wrap_parameters_to_optimization_configs(asymmetry: bool,
per_channel: bool
) -> OptimizationConfigs:
"""
Given the list of model optimization parameters, create OptimizationConfigs data structure.
:param asymmetry: bool. Whether to use asymmetry in quantization.
:param per_channel: bool. Whether to use per channel quantization.
:return: OptimizationConfigs.
"""
quantization_configs = create_quantization_configs(asymmetry=asymmetry, per_channel=per_channel)
calibration_configs = api_calibration_configs()
return OptimizationConfigs(calibration_configs=calibration_configs,
quantization_configs=quantization_configs)