307 lines
17 KiB
Python
307 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
from utils import *
|
|
from argparse import ArgumentParser
|
|
import os
|
|
import sys
|
|
from quantize_types import QuantizerType
|
|
|
|
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):
|
|
nn = VSInn()
|
|
net = nn.create_net()
|
|
|
|
model = model_filename + ".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)
|
|
|
|
return net
|
|
|
|
#set the quantize params got input or output
|
|
def set_input_output_quant_params(quantized_net, layer_lid, in_out_quantized):
|
|
nn = VSInn()
|
|
inputs_outputs = []
|
|
for lid in quantized_net.get_layers():
|
|
l = quantized_net.get_layer(lid)
|
|
if l.is_op("input"):
|
|
inputs_outputs.append(lid)
|
|
elif l.is_op("output"):
|
|
for l_in_lid in l.get_input_layers():
|
|
inputs_outputs.append(l_in_lid)
|
|
|
|
layer_lids = layer_lid.split(',')
|
|
for lid in layer_lids:
|
|
if lid in inputs_outputs:
|
|
layer = quantized_net.get_layer(layer_lid)
|
|
outputs = layer.get_outputs()
|
|
for output in outputs:
|
|
if output.quant_param == None:
|
|
# set the quantize params for activation
|
|
port_name = output.url.split(":")[-1]
|
|
nn.set_quant_params(quantized_net, layer.lid + ":" + port_name,
|
|
{"quantizer": in_out_quantized[0],
|
|
"qtype": in_out_quantized[1]})
|
|
else:
|
|
print("The layer {} has been quantified, so {} is not used for quantization.".format(lid, in_out_quantized))
|
|
else:
|
|
print("Please enter the correct layer name of model inputs or model outputs in json file.")
|
|
return quantized_net
|
|
|
|
# set the quantize params for activation or weight
|
|
def set_quant_params(quantized_net, activation_quantizer=None, weight_quantizer=None):
|
|
nn = VSInn()
|
|
for url, tensor in quantized_net.get_tensors().items():
|
|
if tensor.quant_param is not None:
|
|
layer = quantized_net.get_layer_by_url(url)
|
|
port_name = url.split(":")[-1]
|
|
if weight_quantizer is not None:
|
|
if port_name == "weight":
|
|
# set the quantize params for weight
|
|
nn.set_quant_params(quantized_net, layer.lid + ":" + port_name,
|
|
{"quantizer": weight_quantizer[0],
|
|
"qtype": weight_quantizer[1]})
|
|
if activation_quantizer is not None:
|
|
if port_name != "bias" and port_name != "weight":
|
|
# set the quantize params for activation
|
|
nn.set_quant_params(quantized_net, layer.lid + ":" + port_name,
|
|
{"quantizer": activation_quantizer[0],
|
|
"qtype": activation_quantizer[1]})
|
|
return quantized_net
|
|
|
|
# remove the quant_param of where op when where is followed by softmax
|
|
def remove_quant_params(quantized_net):
|
|
for url, tensor in quantized_net.get_tensors().items():
|
|
layer = quantized_net.get_layer_by_url(url)
|
|
if layer.is_op("softmax"):
|
|
input = layer.get_inputs()
|
|
for t in input:
|
|
layer = quantized_net.get_layer_by_url(t.url)
|
|
if layer.is_op("where"):
|
|
outputs = layer.get_outputs()
|
|
if len(outputs) == 1:
|
|
outputs[0].quant_param = None
|
|
return quantized_net
|
|
|
|
def quantize(net, model_filename, quantized='asymu8', algorithm=1, iterations=1,
|
|
compute_entropy=False, minimize_layer_error=False, layer_lid=None, in_out_quantized=None, save=True):
|
|
|
|
if minimize_layer_error:
|
|
quantized_output = model_filename + '_' + quantized + '.mle.quantize'
|
|
else:
|
|
quantized_output = model_filename + '_' + quantized + '.quantize'
|
|
if os.path.exists(quantized_output) is True:
|
|
print("Delete the {}".format(quantized_output))
|
|
os.system("rm -rf {}".format(quantized_output))
|
|
|
|
nn = VSInn()
|
|
nn.set_device(device='CPU')
|
|
algorithms = ["normal", "kl_divergence", "moving_average", "auto"]
|
|
if quantized in ['e5m2pcqf8', 'e4m3pcqf8', 'e5m2fp8', 'e4m3fp8']:
|
|
algorithm = 0
|
|
print("Your quantization format is '{}', forced use algorithm '{}'.".format(quantized, algorithms[0]))
|
|
|
|
# start quantize
|
|
quantized_format = QuantizerType.get_options()
|
|
if quantized in quantized_format:
|
|
default_support_quantizer_dict = QuantizerType.get_default_support_quantizer_dict()
|
|
a_w_diff_quantizer_dict = QuantizerType.get_a_w_diff_quantizer_dict()
|
|
a_w_same_quantizer_dict = QuantizerType.get_a_w_same_quantizer_dict()
|
|
if in_out_quantized in default_support_quantizer_dict:
|
|
in_out_quantized_dict = default_support_quantizer_dict[in_out_quantized]
|
|
elif in_out_quantized in a_w_same_quantizer_dict:
|
|
in_out_quantized_dict = a_w_same_quantizer_dict[in_out_quantized]
|
|
|
|
if quantized in default_support_quantizer_dict.keys():
|
|
quantizer_dict = default_support_quantizer_dict[quantized]
|
|
print_params(nn.quantize, model=model_filename + ".json", data=model_filename + ".data",
|
|
quantize=model_filename + '_' + quantized + '.quantize', with_input_meta=model_filename + "_inputmeta.yml",
|
|
quantizer=quantizer_dict[0], qtype=quantizer_dict[1], algorithm=algorithms[algorithm], iterations=iterations, rebuild=True,
|
|
compute_entropy=compute_entropy, minimize_layer_error=minimize_layer_error)
|
|
quantized_net = nn.quantize(net, quantizer=quantizer_dict[0], qtype=quantizer_dict[1],
|
|
algorithm=algorithms[algorithm], iterations=iterations, rebuild=True,
|
|
compute_entropy=compute_entropy, minimize_layer_error=minimize_layer_error, divergence_first_quantize_bits=13)
|
|
if layer_lid is not None and in_out_quantized is not None:
|
|
quantized_net = set_input_output_quant_params(quantized_net, layer_lid, in_out_quantized_dict)
|
|
quantized_net = nn.quantize(quantized_net, quantizer=quantizer_dict[0], qtype=quantizer_dict[1],
|
|
algorithm=algorithms[algorithm], iterations=iterations, rebuild=False,
|
|
compute_entropy=compute_entropy, minimize_layer_error=minimize_layer_error,
|
|
divergence_first_quantize_bits=13)
|
|
|
|
elif quantized in a_w_diff_quantizer_dict.keys():
|
|
quantizer_dict = list(a_w_diff_quantizer_dict[quantized].values())
|
|
activation_quantizer_dict = quantizer_dict[0]
|
|
weight_quantizer_dict = quantizer_dict[1]
|
|
|
|
# The first quantization to get the quantize_tab and quantize weight,
|
|
# if you want to quantize all layers, set rebuild_all = True
|
|
print_params(nn.quantize, model=model_filename + ".json", data=model_filename + ".data",
|
|
quantize=model_filename + '_' + quantized + '.quantize', with_input_meta=model_filename + "_inputmeta.yml",
|
|
quantizer=[activation_quantizer_dict[0], weight_quantizer_dict[0]],
|
|
qtype=[activation_quantizer_dict[1], weight_quantizer_dict[1]], algorithm=algorithms[algorithm],
|
|
iterations=1, compute_entropy=compute_entropy, minimize_layer_error=minimize_layer_error)
|
|
if 'symi16' in a_w_diff_quantizer_dict[quantized].keys() or "asymu16" in a_w_diff_quantizer_dict[quantized].keys() \
|
|
or 'float16' in a_w_diff_quantizer_dict[quantized].keys() :
|
|
quantized_net = nn.quantize(net, quantizer=weight_quantizer_dict[0], qtype=weight_quantizer_dict[1],
|
|
algorithm=algorithms[algorithm], iterations=iterations, rebuild=True, divergence_first_quantize_bits=13)
|
|
# set the quantize params for specified ops that you want.
|
|
# Here is set the activation op quantize params.
|
|
quantized_net = set_quant_params(quantized_net, activation_quantizer=activation_quantizer_dict)
|
|
if layer_lid is not None and in_out_quantized is not None:
|
|
quantized_net = set_input_output_quant_params(quantized_net, layer_lid, in_out_quantized_dict)
|
|
|
|
# The second quantization takes effect on the quantize params to quantize activation.
|
|
# The rebuild must be False
|
|
quantized_net = nn.quantize(quantized_net, quantizer=weight_quantizer_dict[0],
|
|
qtype=weight_quantizer_dict[1], algorithm=algorithms[algorithm], iterations=iterations,
|
|
rebuild=False, compute_entropy=compute_entropy,
|
|
minimize_layer_error=minimize_layer_error, divergence_first_quantize_bits=13)
|
|
else:
|
|
quantized_net = nn.quantize(net, quantizer=activation_quantizer_dict[0], qtype=activation_quantizer_dict[1],
|
|
algorithm=algorithms[algorithm], iterations=iterations, rebuild=True,
|
|
divergence_first_quantize_bits=13)
|
|
# set the quantize params for specified ops that you want.
|
|
# Here is set the weight op quantize params.
|
|
quantized_net = set_quant_params(quantized_net, weight_quantizer=weight_quantizer_dict)
|
|
if layer_lid is not None and in_out_quantized is not None:
|
|
quantized_net = set_input_output_quant_params(quantized_net, layer_lid, in_out_quantized_dict)
|
|
|
|
# The second quantization takes effect on the quantize params to quantize weight.
|
|
# The rebuild must be False
|
|
quantized_net = nn.quantize(quantized_net, quantizer=activation_quantizer_dict[0],
|
|
qtype=activation_quantizer_dict[1], algorithm=algorithms[algorithm],
|
|
iterations=iterations,
|
|
rebuild=False, compute_entropy=compute_entropy,
|
|
minimize_layer_error=minimize_layer_error,
|
|
divergence_first_quantize_bits=13)
|
|
elif quantized in a_w_same_quantizer_dict.keys():
|
|
quantizer_dict = a_w_same_quantizer_dict[quantized]
|
|
|
|
# The first quantization to get the quantize_tab,
|
|
# if you want to quantize all layers, set rebuild_all = True
|
|
print_params(nn.quantize, model=model_filename + ".json", data=model_filename + ".data",
|
|
quantize=model_filename + '_' + quantized + '.quantize', with_input_meta=model_filename + "_inputmeta.yml",
|
|
quantizer=quantizer_dict[0], qtype=quantizer_dict[1], algorithm=algorithms[algorithm], iterations=1,
|
|
compute_entropy=compute_entropy, minimize_layer_error=minimize_layer_error)
|
|
quantized_net = nn.quantize(net, quantizer="symmetric_affine", qtype="int8",
|
|
algorithm=algorithms[algorithm], iterations=1, rebuild=True, divergence_first_quantize_bits=13)
|
|
# set the quantize params for specified ops
|
|
quantized_net = set_quant_params(quantized_net, activation_quantizer=quantizer_dict, weight_quantizer=quantizer_dict)
|
|
if layer_lid is not None and in_out_quantized is not None:
|
|
quantized_net = set_input_output_quant_params(quantized_net, layer_lid, in_out_quantized_dict)
|
|
|
|
# The second quantization takes effect on the quantize params.
|
|
# The rebuild must be False
|
|
quantized_net = nn.quantize(quantized_net, quantizer="symmetric_affine", qtype="int8",
|
|
algorithm=algorithms[algorithm], iterations=iterations, rebuild=False,
|
|
compute_entropy=compute_entropy, minimize_layer_error=minimize_layer_error, divergence_first_quantize_bits=13)
|
|
quantized_net = remove_quant_params(quantized_net)
|
|
|
|
if save:
|
|
nn.save_model_quantize(quantized_net, quantized_output)
|
|
if algorithm == 3:
|
|
output_model = model_filename + '_auto.json'
|
|
nn.save_model(quantized_net, output_model)
|
|
print("You use the \'auto\' algorithm, then the {}_auto.json file has been generated!".format(model_filename))
|
|
if minimize_layer_error:
|
|
output_data = model_filename + '_mle.data'
|
|
nn.save_model_data(quantized_net, output_data)
|
|
return quantized_net
|
|
else:
|
|
print("Please enter the correct quantization format.")
|
|
print(list(quantized_format))
|
|
sys.exit(1)
|
|
|
|
def qat_quantize(net, layer_lid, in_out_quantized, quantize_file):
|
|
nn = VSInn()
|
|
if os.path.exists(quantize_file) is True:
|
|
nn.load_model_quantize(net, quantize_file)
|
|
else:
|
|
print("The quantize file {} does not exist.")
|
|
sys.exit(1)
|
|
|
|
default_support_quantizer_dict = QuantizerType.get_default_support_quantizer_dict()
|
|
a_w_same_quantizer_dict = QuantizerType.get_a_w_same_quantizer_dict()
|
|
if in_out_quantized in default_support_quantizer_dict:
|
|
in_out_quantized_dict = default_support_quantizer_dict[in_out_quantized]
|
|
elif in_out_quantized in a_w_same_quantizer_dict:
|
|
in_out_quantized_dict = a_w_same_quantizer_dict[in_out_quantized]
|
|
|
|
if layer_lid is not None and in_out_quantized is not None:
|
|
quantized_net = set_input_output_quant_params(net, layer_lid, in_out_quantized_dict)
|
|
quantized_net = nn.quantize(quantized_net, quantizer=in_out_quantized_dict[0], qtype=in_out_quantized_dict[1], rebuild=False)
|
|
quantized_output = quantize_file.split('.quantize')[0] + "_in_out.quantize"
|
|
nn.save_model_quantize(quantized_net, quantized_output)
|
|
|
|
def main():
|
|
options = ArgumentParser()
|
|
options.add_argument("model", type=str, help="Model directory")
|
|
options.add_argument("quantized", type=str, help="Quantization type. Including " + ', '.join(list(QuantizerType.get_options())))
|
|
options.add_argument("--algorithm", type=int, help="Quantization algotithm. The corresponding relationship between numbers and algorithms is as follows:"
|
|
"[0: normal, 1:kl_divergence, 2:moving_average, 3:auto])",
|
|
default=1, choices=[0, 1, 2, 3])
|
|
options.add_argument("--iterations", type=int, help="Running iterations.", default=1)
|
|
options.add_argument("--entropy", action="store_true", help="Compute tensor entropy.")
|
|
options.add_argument("--mle", action="store_true", help="Minimize per layer error")
|
|
options.add_argument("--lid", type=str, help="The layer names of the model input or model output in json file."
|
|
"The layer names of the same subgraph are separated with commas.")
|
|
options.add_argument("--in_out_quantized", type=str, help="The quantization type of the model input or model output in json file.")
|
|
options.add_argument("--is_qat", action="store_true", help="Whether the model is QAT model.")
|
|
options.add_argument("--quantize_file", type=str, help="If model is the QAT model, please specify the path of the quantize file.")
|
|
|
|
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
|
|
algorithm = args.algorithm
|
|
iterations = args.iterations
|
|
compute_entropy = args.entropy
|
|
minimize_layer_error = args.mle
|
|
lid = args.lid
|
|
in_out_quantized = args.in_out_quantized
|
|
is_qat = args.is_qat
|
|
quantize_file = args.quantize_file
|
|
|
|
#load_net
|
|
net = load_net(model_filename)
|
|
|
|
#quantize
|
|
if is_qat:
|
|
qat_quantize(net, lid, in_out_quantized, quantize_file)
|
|
else:
|
|
quantize(net, model_filename, quantized, algorithm, iterations, compute_entropy, minimize_layer_error, lid, in_out_quantized)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|