netrans/script/dump.py

114 lines
4.0 KiB
Python
Executable File

#!/usr/bin/env python3
from argparse import ArgumentParser
import os
import sys
from quantize_types import QuantizerType
from utils import *
import importlib
try:
importlib.import_module("acuitylib")
except:
ACUITY_PATH = os.environ['ACUITY_PATH']
sys.path.append(ACUITY_PATH)
from acuitylib.vsi_nn import VSInn
def load_net(model_filename, quantized, use_hybrid=False):
nn = VSInn()
net = nn.create_net()
if not use_hybrid:
model = model_filename + ".json"
else:
model = model_filename + "_" + quantized + "_hy.quantize.json"
data = model_filename + ".data"
inputmeta = model_filename + "_inputmeta.yml"
if os.path.exists(model) is True:
nn.load_model(net, model)
else:
print("{} file does not exists.".format(model))
sys.exit(1)
if os.path.exists(data) is True:
nn.load_model_data(net, data)
else:
print("{} file does not exists.".format(data))
sys.exit(1)
if os.path.exists(inputmeta) is True:
nn.load_model_inputmeta(net, inputmeta)
else:
print("{} file does not exists.".format(inputmeta))
sys.exit(1)
if quantized != "float32":
if not use_hybrid:
model_quantize = model_filename + '_' + quantized + ".quantize"
else:
model_quantize = model_filename + '_' + quantized + "_hy.quantize"
if os.path.exists(model_quantize) is True:
nn.load_model_quantize(net, model_quantize)
else:
print('{} does not exist'.format(model_quantize))
sys.exit(1)
return net
def dump(net, model_filename, quantized='asymu8', use_hybrid=False):
nn = VSInn()
if not use_hybrid:
quantize_file = model_filename + '_' + quantized + ".quantize"
model = model_filename + ".json"
output_dir = 'dump/{}_{}/'.format(model_filename, quantized)
else:
# add hybrid quantize for print log
quantize_file = model_filename + '_' + quantized + "_hy.quantize"
model = model_filename + '_' + quantized + "_hy.quantize.json"
# add hybrid quantize output file name
output_dir = 'dump/{}_{}/'.format(model_filename, quantized + "_hy")
if quantized != "float32":
print_params(nn.dump, model=model, data=model_filename + ".data", quantize=quantize_file,
with_input_meta=model_filename + "_inputmeta.yml", output_path=output_dir)
else:
print_params(nn.dump, model=model, data=model_filename + ".data",
with_input_meta=model_filename + "_inputmeta.yml", output_path=output_dir)
nn.dump(net, output_path=output_dir)
def main():
options = ArgumentParser()
options.add_argument("model", type=str, help="Model directory")
options.add_argument("quantized", type=str, help="Quantization type, including float32, " + ', '.join(list(QuantizerType.get_options()))
+ ", \'float32\' means not quantized.")
options.add_argument("--use_hybrid", action="store_true",
help="if you use hybrid quantize,please set this --use_hybrid")
args = options.parse_args()
print(args)
if os.path.exists(args.model) and os.path.isdir(os.path.abspath(args.model)):
model_filename = get_modelfile_name(args.model)
if model_filename is None:
print("Please enter the path that includes the model.")
os.chdir(args.model)
else:
model_filename = args.model
quantized = args.quantized
use_hybrid = args.use_hybrid
quantized_format = QuantizerType.get_options()
if quantized not in quantized_format and quantized != 'float32':
print("Please enter the correct quantization format.")
quantized_format.insert(0, 'float32')
print(list(quantized_format))
sys.exit(1)
# load net
net = load_net(model_filename, quantized, use_hybrid)
#dump
dump(net, model_filename, quantized, use_hybrid)
if __name__ == "__main__":
main()