Function main

tools/gen_header_for_bin_reduce.py:261–301

()

def main():
    parser = argparse.ArgumentParser(
        description="generate header file for reducing binary size by "
        "stripping unused oprs in a particular network; output file would "
        "be written to bin_reduce.h",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "inputs",
        nargs="+",
        help="input files that describe specific traits of the network; "
        "can be one of the following:"
        "  1. json files generated by "
        "megbrain.serialize_comp_graph_to_file() in python; "
        "  2. trace files generated by midout library",
    )
    default_file = os.path.join(
        HeaderGen.get_megengine_root(), "src", "bin_reduce_cmake.h"
    )
    is_megvii3 = HeaderGen.get_megvii3_root()
    if is_megvii3:
        default_file = os.path.join(
            HeaderGen.get_megvii3_root(), "utils", "bin_reduce.h"
        )
    parser.add_argument("-o", "--output", help="output file", default=default_file)
    args = parser.parse_args()
    print("config output file: {}".format(args.output))

    gen = HeaderGen()
    for i in args.inputs:
        print("==== processing {}".format(i))
        with open(i) as fin:
            if fin.read(len(MIDOUT_TRACE_MAGIC)) == MIDOUT_TRACE_MAGIC:
                gen.extend_midout(i)
                gen.extend_elemwise_mode_info(i)
            else:
                fin.seek(0)
                gen.extend_netinfo(json.loads(fin.read()))

    with open(args.output, "w") as fout:
        gen.generate(fout)


if __name__ == "__main__":