#! /usr/bin/env python3
"""
Load a YAML file and use the JSON and npz file that are referred in the YAML
to load the pre-calibrated AwesomeNet, quantize it, and generate MLC files.
"""
import argparse
import afe
from afe.core.utils import load_configs_from_yaml as _load_configs_from_yaml
def _parse_args():
parser = argparse.ArgumentParser(description='Load a calibrated model from YAML file.'
'Quantize it and generate the MLC files')
parser.add_argument('-yaml', '--yaml', required=True, type=str,
help='YAML file that has model configuration, optimization configuration,'
' and JSON and npz file path to load the quantized model')
args = parser.parse_args()
return args
def _main(yaml_file):
# Load model config and optimization config from a YAML file
model_config, opt_config = _load_configs_from_yaml(yaml_file)
# Load calibrated model
target_net_file_name = model_config.name + afe.CALIBRATED_POSTFIX
net = afe.load_awesomenet(model_name=target_net_file_name,
network_directory=model_config.output_directory)
# Quantize the model
afe.quantize_network(net=net,
model_config=model_config,
opt_config=opt_config,
input_generator=None)
afe.compile_network(net=net,
model_config=model_config,
opt_config=opt_config,
enable_large_tensors=True)
[docs]
def main():
args = _parse_args()
_main(args.yaml)
if __name__ == "__main__":
main()