Update infer_rec_prof.py

tools/infer_rec_prof.py

+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
 
+import numpy as np
+
+import os
+import sys
+import json
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(__dir__)
+sys.path.insert(0, os.path.abspath(os.path.join(__dir__, "..")))
+
+os.environ["FLAGS_allocator_strategy"] = "auto_growth"
+
+import paddle
+import paddle.profiler as profiler  # 导入性能分析器模块
+
+from ppocr.data import create_operators, transform
+from ppocr.modeling.architectures import build_model
+from ppocr.postprocess import build_post_process
+from ppocr.utils.save_load import load_model
+from ppocr.utils.utility import get_image_file_list
+import tools.program as program
+
+
+def main():
+    global_config = config["Global"]
+
+    # build post process
+    post_process_class = build_post_process(config["PostProcess"], global_config)
+
+    # build model
+    if hasattr(post_process_class, "character"):
+        char_num = len(getattr(post_process_class, "character"))
+        if config["Architecture"]["algorithm"] in [
+            "Distillation",
+        ]:  # distillation model
+            for key in config["Architecture"]["Models"]:
+                if (
+                    config["Architecture"]["Models"][key]["Head"]["name"] == "MultiHead"
+                ):  # multi head
+                    out_channels_list = {}
+                    if config["PostProcess"]["name"] == "DistillationSARLabelDecode":
+                        char_num = char_num - 2
+                    if config["PostProcess"]["name"] == "DistillationNRTRLabelDecode":
+                        char_num = char_num - 3
+                    out_channels_list["CTCLabelDecode"] = char_num
+                    out_channels_list["SARLabelDecode"] = char_num + 2
+                    out_channels_list["NRTRLabelDecode"] = char_num + 3
+                    config["Architecture"]["Models"][key]["Head"][
+                        "out_channels_list"
+                    ] = out_channels_list
+                else:
+                    config["Architecture"]["Models"][key]["Head"][
+                        "out_channels"
+                    ] = char_num
+        elif config["Architecture"]["Head"]["name"] == "MultiHead":  # multi head
+            out_channels_list = {}
+            char_num = len(getattr(post_process_class, "character"))
+            if config["PostProcess"]["name"] == "SARLabelDecode":
+                char_num = char_num - 2
+            if config["PostProcess"]["name"] == "NRTRLabelDecode":
+                char_num = char_num - 3
+            out_channels_list["CTCLabelDecode"] = char_num
+            out_channels_list["SARLabelDecode"] = char_num + 2
+            out_channels_list["NRTRLabelDecode"] = char_num + 3
+            config["Architecture"]["Head"]["out_channels_list"] = out_channels_list
+        else:  # base rec model
+            config["Architecture"]["Head"]["out_channels"] = char_num
+
+    if config["Architecture"].get("algorithm") in ["LaTeXOCR"]:
+        config["Architecture"]["Backbone"]["is_predict"] = True
+        config["Architecture"]["Backbone"]["is_export"] = True
+        config["Architecture"]["Head"]["is_export"] = True
+
+    model = build_model(config["Architecture"])
+
+    load_model(config, model)
+
+    # create data ops
+    transforms = []
+    for op in config["Eval"]["dataset"]["transforms"]:
+        op_name = list(op)[0]
+        if "Label" in op_name:
+            continue
+        elif op_name in ["RecResizeImg"]:
+            op[op_name]["infer_mode"] = True
+        elif op_name == "KeepKeys":
+            if config["Architecture"]["algorithm"] == "SRN":
+                op[op_name]["keep_keys"] = [
+                    "image",
+                    "encoder_word_pos",
+                    "gsrm_word_pos",
+                    "gsrm_slf_attn_bias1",
+                    "gsrm_slf_attn_bias2",
+                ]
+            elif config["Architecture"]["algorithm"] == "SAR":
+                op[op_name]["keep_keys"] = ["image", "valid_ratio"]
+            elif config["Architecture"]["algorithm"] == "RobustScanner":
+                op[op_name]["keep_keys"] = ["image", "valid_ratio", "word_positons"]
+            else:
+                op[op_name]["keep_keys"] = ["image"]
+        transforms.append(op)
+    global_config["infer_mode"] = True
+    ops = create_operators(transforms, global_config)
+
+    save_res_path = config["Global"].get(
+        "save_res_path", "./output/rec/predicts_rec.txt"
+    )
+    if not os.path.exists(os.path.dirname(save_res_path)):
+        os.makedirs(os.path.dirname(save_res_path))
+
+    model.eval()
+
+    # 创建性能分析器相关的代码
+    def my_on_trace_ready(prof):
+        callback = profiler.export_chrome_tracing('./profiler_demo')
+        callback(prof)
+
+        # 将 Overview Summary 和 Operator Summary 保存到文件
+        with open('./profiler_summary.txt', 'w') as f:
+            f.write("Overview Summary:\n")
+            summary_overview = prof.summary(sorted_by=profiler.SortedKeys.GPUTotal,
+                                            op_detail=True,
+                                            thread_sep=True,
+                                            time_unit='ms')
+            if summary_overview is not None:
+                f.write(summary_overview)
+            else:
+                f.write("No summary available for Overview.\n")
+
+            f.write("\n\nOperator Summary:\n")
+            summary_operator = prof.summary(sorted_by=profiler.SortedKeys.GPUTotal,
+                                            op_detail=True,
+                                            thread_sep=True,
+                                            time_unit='ms')
+            if summary_operator is not None:
+                f.write(summary_operator)
+            else:
+                f.write("No summary available for Operator.\n")
+
+    # 初始化 Profiler 对象，设置 timer_only=False 以收集详细信息
+    p = profiler.Profiler(on_trace_ready=my_on_trace_ready,
+                          timer_only=False)
+    p.start()
+
+    infer_imgs = config["Global"]["infer_img"]
+    infer_list = config["Global"].get("infer_list", None)
+    with open(save_res_path, "w") as fout:
+        for file in get_image_file_list(infer_imgs, infer_list=infer_list):
+            logger.info("infer_img: {}".format(file))
+            with open(file, "rb") as f:
+                img = f.read()
+                if config["Architecture"]["algorithm"] in [
+                    "UniMERNet",
+                    "PP-FormulaNet-S",
+                    "PP-FormulaNet-L",
+                ]:
+                    data = {"image": img, "filename": file}
+                else:
+                    data = {"image": img}
+            batch = transform(data, ops)
+            if config["Architecture"]["algorithm"] == "SRN":
+                encoder_word_pos_list = np.expand_dims(batch[1], axis=0)
+                gsrm_word_pos_list = np.expand_dims(batch[2], axis=0)
+                gsrm_slf_attn_bias1_list = np.expand_dims(batch[3], axis=0)
+                gsrm_slf_attn_bias2_list = np.expand_dims(batch[4], axis=0)
+
+                others = [
+                    paddle.to_tensor(encoder_word_pos_list),
+                    paddle.to_tensor(gsrm_word_pos_list),
+                    paddle.to_tensor(gsrm_slf_attn_bias1_list),
+                    paddle.to_tensor(gsrm_slf_attn_bias2_list),
+                ]
+            if config["Architecture"]["algorithm"] == "SAR":
+                valid_ratio = np.expand_dims(batch[-1], axis=0)
+                img_metas = [paddle.to_tensor(valid_ratio)]
+            if config["Architecture"]["algorithm"] == "RobustScanner":
+                valid_ratio = np.expand_dims(batch[1], axis=0)
+                word_positons = np.expand_dims(batch[2], axis=0)
+                img_metas = [
+                    paddle.to_tensor(valid_ratio),
+                    paddle.to_tensor(word_positons),
+                ]
+            if config["Architecture"]["algorithm"] == "CAN":
+                image_mask = paddle.ones(
+                    (np.expand_dims(batch[0], axis=0).shape), dtype="float32"
+                )
+                label = paddle.ones((1, 36), dtype="int64")
+            images = np.expand_dims(batch[0], axis=0)
+            images = paddle.to_tensor(images)
+            if config["Architecture"]["algorithm"] == "SRN":
+                preds = model(images, others)
+            elif config["Architecture"]["algorithm"] == "SAR":
+                preds = model(images, img_metas)
+            elif config["Architecture"]["algorithm"] == "RobustScanner":
+                preds = model(images, img_metas)
+            elif config["Architecture"]["algorithm"] == "CAN":
+                preds = model([images, image_mask, label])
+            else:
+                preds = model(images)
+            post_result = post_process_class(preds)
+            info = None
+            if isinstance(post_result, dict):
+                rec_info = dict()
+                for key in post_result:
+                    if len(post_result[key][0]) >= 2:
+                        rec_info[key] = {
+                            "label": post_result[key][0][0],
+                            "score": float(post_result[key][0][1]),
+                        }
+                info = json.dumps(rec_info, ensure_ascii=False)
+            elif isinstance(post_result, list) and isinstance(post_result[0], int):
+                # for RFLearning CNT branch
+                info = str(post_result[0])
+            elif config["Architecture"]["algorithm"] in [
+                "LaTeXOCR",
+                "UniMERNet",
+                "PP-FormulaNet-S",
+                "PP-FormulaNet-L",
+            ]:
+                info = str(post_result[0])
+            else:
+                if len(post_result[0]) >= 2:
+                    info = post_result[0][0] + "\t" + str(post_result[0][1])
+
+            if info is not None:
+                logger.info("\t result: {}".format(info))
+                fout.write(file + "\t" + info + "\n")
+            p.step()  # 每次推理后调用 profiler 的 step 方法
+    p.stop()  # 停止 profiler
+    logger.info("success!")
+
+
+if __name__ == "__main__":
+    config, device, logger, vdl_writer = program.preprocess()
+    main()