76 lines
3.5 KiB
Python
76 lines
3.5 KiB
Python
import json
|
|
import yaml
|
|
import sys
|
|
import os
|
|
|
|
|
|
def load_config(config_path):
|
|
with open(config_path, "r") as file:
|
|
config = yaml.safe_load(file)
|
|
|
|
if "quant" in config:
|
|
quant_info = config["quant"]
|
|
if "output_path" in quant_info:
|
|
output_path = os.path.join(quant_info["output_path"], config["task_id"])
|
|
if "march" in quant_info:
|
|
march = "nash" if "nash" in quant_info["march"] else "bayes"
|
|
convert_yaml_path = f"pyq_yaml/{march}/"
|
|
# prepare the nash and bayes bpu
|
|
## first prepare the VisionEncoder yaml
|
|
VisionEncoder_yaml_path = os.path.join(convert_yaml_path, "VisionEncoder.yaml")
|
|
with open(VisionEncoder_yaml_path, "r") as file:
|
|
VisionEncoder_yaml = yaml.safe_load(file)
|
|
VisionEncoder_yaml["model_parameters"]["onnx_model"] = quant_info["VisionEncoder"]["onnx_model"]
|
|
VisionEncoder_yaml["calibration_parameters"]["cal_data_dir"] = quant_info["VisionEncoder"]["calibration_data"]
|
|
VisionEncoder_yaml["model_parameters"]["march"] = quant_info["march"]
|
|
|
|
# Make sure output ptq_yaml directory exists
|
|
output_ptq_yaml_dir = os.path.join(output_path, "ptq_yaml")
|
|
os.makedirs(output_ptq_yaml_dir, exist_ok=True)
|
|
|
|
# Save VisionEncoder yaml to output/ptq_yaml
|
|
VisionEncoder_yaml_save_path = os.path.join(output_ptq_yaml_dir, "VisionEncoder.yaml")
|
|
with open(VisionEncoder_yaml_save_path, "w") as file:
|
|
yaml.safe_dump(VisionEncoder_yaml, file)
|
|
|
|
## second prepare the TransformerLayers yaml
|
|
TransformerLayers_yaml_path = os.path.join(convert_yaml_path, "TransformerLayers.yaml")
|
|
with open(TransformerLayers_yaml_path, "r") as file:
|
|
TransformerLayers_yaml = yaml.safe_load(file)
|
|
TransformerLayers_yaml["model_parameters"]["onnx_model"] = quant_info["TransformerLayers"]["onnx_model"]
|
|
TransformerLayers_yaml["model_parameters"]["march"] = quant_info["march"]
|
|
|
|
TransformerLayers_Cal_dir = quant_info["TransformerLayers"]["calibration_data"]
|
|
# (Fix cal_data_dir variable)
|
|
cal_data_dir = TransformerLayers_Cal_dir
|
|
sub_dirs = [d for d in os.listdir(cal_data_dir) if os.path.isdir(os.path.join(cal_data_dir, d))]
|
|
input_names = []
|
|
for name in sub_dirs:
|
|
if name == "state":
|
|
input_names.append("states")
|
|
else:
|
|
input_names.append(f"{name}_features")
|
|
input_name_str = ";".join(input_names) + ";"
|
|
|
|
TransformerLayers_yaml["input_parameters"]["input_name"] = input_name_str
|
|
TransformerLayers_yaml["input_parameters"]["input_type_rt"] = "featuremap;" * len(input_names)
|
|
TransformerLayers_yaml["input_parameters"]["input_layout_rt"] = "NCHW;" * len(input_names)
|
|
TransformerLayers_yaml["input_parameters"]["input_type_train"] = "featuremap;" * len(input_names)
|
|
TransformerLayers_yaml["input_parameters"]["input_layout_train"] = "NCHW;" * len(input_names)
|
|
TransformerLayers_yaml["input_parameters"]["norm_type"] = "no_preprocess;" * len(input_names)
|
|
TransformerLayers_yaml["calibration_parameters"]["cal_data_dir"] = ";".join([os.path.join(TransformerLayers_Cal_dir, name) for name in sub_dirs]) + ";"
|
|
TransformerLayers_yaml["calibration_parameters"]["cal_data_type"] = "float32;" * len(input_names)
|
|
|
|
# Save TransformerLayers yaml to output/ptq_yaml
|
|
TransformerLayers_yaml_save_path = os.path.join(output_ptq_yaml_dir, "TransformerLayers.yaml")
|
|
with open(TransformerLayers_yaml_save_path, "w") as file:
|
|
yaml.safe_dump(TransformerLayers_yaml, file)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
config_path = sys.argv[1]
|
|
config = load_config(config_path) |