Merge remote-tracking branch 'origin/dygraph' into dygraph

ppstructure/utility.py

@@ -21,13 +21,31 @@ def init_args():
     parser = infer_args()
 
     # params for output
-    parser.add_argument("--output", type=str, default='./output/table')
+    parser.add_argument("--output", type=str, default='./output')
     # params for table structure
     parser.add_argument("--table_max_len", type=int, default=488)
     parser.add_argument("--table_model_dir", type=str)
     parser.add_argument("--table_char_type", type=str, default='en')
-    parser.add_argument("--table_char_dict_path", type=str, default="../ppocr/utils/dict/table_structure_dict.txt")
-    parser.add_argument("--layout_path_model", type=str, default="lp://PubLayNet/ppyolov2_r50vd_dcn_365e_publaynet/config")
+    parser.add_argument(
+        "--table_char_dict_path",
+        type=str,
+        default="../ppocr/utils/dict/table_structure_dict.txt")
+    parser.add_argument(
+        "--layout_path_model",
+        type=str,
+        default="lp://PubLayNet/ppyolov2_r50vd_dcn_365e_publaynet/config")
+
+    # params for ser
+    parser.add_argument("--model_name_or_path", type=str)
+    parser.add_argument("--max_seq_length", type=int, default=512)
+    parser.add_argument(
+        "--label_map_path", type=str, default='./vqa/labels/labels_ser.txt')
+
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default='structure',
+        help='structure and vqa is supported')
     return parser
 
 
@@ -48,5 +66,6 @@ def draw_structure_result(image, result, font_path):
                 boxes.append(np.array(box).reshape(-1, 2))
                 txts.append(rec_res[0])
                 scores.append(rec_res[1])
-    im_show = draw_ocr_box_txt(image, boxes, txts, scores, font_path=font_path,drop_score=0)
+    im_show = draw_ocr_box_txt(
+        image, boxes, txts, scores, font_path=font_path, drop_score=0)
     return im_show

ppstructure/vqa/README.md

@@ -18,12 +18,13 @@ PP-Structure 里的 DOC-VQA算法基于PaddleNLP自然语言处理算法库进
 
 ## 1 性能
 
-我们在 [XFUN](https://github.com/doc-analysis/XFUND) 评估数据集上对算法进行了评估，性能如下
+我们在 [XFUN](https://github.com/doc-analysis/XFUND) 的中文数据集上对算法进行了评估，性能如下
 
-|任务|    f1 | 模型下载地址|
-|:---:|:---:| :---:|
-|SER|0.9056| [链接](https://paddleocr.bj.bcebos.com/pplayout/PP-Layout_v1.0_ser_pretrained.tar)|
-|RE|0.7113| [链接](https://paddleocr.bj.bcebos.com/pplayout/PP-Layout_v1.0_re_pretrained.tar)|
+| 模型 | 任务 | f1 | 模型下载地址 |
+|:---:|:---:|:---:| :---:|
+| LayoutXLM | RE | 0.7113 | [链接](https://paddleocr.bj.bcebos.com/pplayout/PP-Layout_v1.0_re_pretrained.tar) |
+| LayoutXLM | SER | 0.9056 | [链接](https://paddleocr.bj.bcebos.com/pplayout/PP-Layout_v1.0_ser_pretrained.tar) |
+| LayoutLM | SER | 0.78 | [链接](https://paddleocr.bj.bcebos.com/pplayout/LayoutLM_ser_pretrained.tar) |
 
 
 
@@ -98,7 +99,7 @@ git clone https://gitee.com/paddlepaddle/PaddleOCR
 # 需要使用PaddleNLP最新的代码版本进行安装
 git clone https://github.com/PaddlePaddle/PaddleNLP -b develop
 cd PaddleNLP
-pip install -e .
+pip3 install -e .
 ```
 
 
@@ -135,13 +136,13 @@ wget https://paddleocr.bj.bcebos.com/dataset/XFUND.tar
 ```shell
 python3.7 train_ser.py \
     --model_name_or_path "layoutxlm-base-uncased" \
+    --ser_model_type "LayoutXLM" \
     --train_data_dir "XFUND/zh_train/image" \
     --train_label_path "XFUND/zh_train/xfun_normalize_train.json" \
     --eval_data_dir "XFUND/zh_val/image" \
     --eval_label_path "XFUND/zh_val/xfun_normalize_val.json" \
     --num_train_epochs 200 \
     --eval_steps 10 \
-    --save_steps 500 \
     --output_dir "./output/ser/" \
     --learning_rate 5e-5 \
     --warmup_steps 50 \
@@ -151,13 +152,50 @@ python3.7 train_ser.py \
 
 最终会打印出`precision`, `recall`, `f1`等指标，模型和训练日志会保存在`./output/ser/`文件夹中。
 
+* 恢复训练
+
+```shell
+python3.7 train_ser.py \
+    --model_name_or_path "model_path" \
+    --ser_model_type "LayoutXLM" \
+    --train_data_dir "XFUND/zh_train/image" \
+    --train_label_path "XFUND/zh_train/xfun_normalize_train.json" \
+    --eval_data_dir "XFUND/zh_val/image" \
+    --eval_label_path "XFUND/zh_val/xfun_normalize_val.json" \
+    --num_train_epochs 200 \
+    --eval_steps 10 \
+    --output_dir "./output/ser/" \
+    --learning_rate 5e-5 \
+    --warmup_steps 50 \
+    --evaluate_during_training \
+    --num_workers 8 \
+    --seed 2048 \
+    --resume
+```
+
+* 评估
+```shell
+export CUDA_VISIBLE_DEVICES=0
+python3 eval_ser.py \
+    --model_name_or_path "PP-Layout_v1.0_ser_pretrained/" \
+    --ser_model_type "LayoutXLM" \
+    --eval_data_dir "XFUND/zh_val/image" \
+    --eval_label_path "XFUND/zh_val/xfun_normalize_val.json" \
+    --per_gpu_eval_batch_size 8 \
+    --num_workers 8 \
+    --output_dir "output/ser/"  \
+    --seed 2048
+```
+最终会打印出`precision`, `recall`, `f1`等指标
+
 * 使用评估集合中提供的OCR识别结果进行预测
 
 ```shell
 export CUDA_VISIBLE_DEVICES=0
 python3.7 infer_ser.py \
-    --model_name_or_path "./PP-Layout_v1.0_ser_pretrained/" \
-    --output_dir "output_res/" \
+    --model_name_or_path "PP-Layout_v1.0_ser_pretrained/" \
+    --ser_model_type "LayoutXLM" \
+    --output_dir "output/ser/" \
     --infer_imgs "XFUND/zh_val/image/" \
     --ocr_json_path "XFUND/zh_val/xfun_normalize_val.json"
 ```
@@ -169,9 +207,10 @@ python3.7 infer_ser.py \
 ```shell
 export CUDA_VISIBLE_DEVICES=0
 python3.7 infer_ser_e2e.py \
-    --model_name_or_path "./output/PP-Layout_v1.0_ser_pretrained/" \
+    --model_name_or_path "PP-Layout_v1.0_ser_pretrained/" \
+    --ser_model_type "LayoutXLM" \
     --max_seq_length 512 \
-    --output_dir "output_res_e2e/" \
+    --output_dir "output/ser_e2e/" \
     --infer_imgs "images/input/zh_val_0.jpg"
 ```
 
@@ -188,6 +227,7 @@ python3.7 helper/eval_with_label_end2end.py --gt_json_path XFUND/zh_val/xfun_nor
 * 启动训练
 
 ```shell
+export CUDA_VISIBLE_DEVICES=0
 python3 train_re.py \
     --model_name_or_path "layoutxlm-base-uncased" \
     --train_data_dir "XFUND/zh_train/image" \
@@ -195,32 +235,74 @@ python3 train_re.py \
     --eval_data_dir "XFUND/zh_val/image" \
     --eval_label_path "XFUND/zh_val/xfun_normalize_val.json" \
     --label_map_path 'labels/labels_ser.txt' \
-    --num_train_epochs 2 \
+    --num_train_epochs 200 \
     --eval_steps 10 \
-    --save_steps 500 \
     --output_dir "output/re/"  \
     --learning_rate 5e-5 \
     --warmup_steps 50 \
     --per_gpu_train_batch_size 8 \
     --per_gpu_eval_batch_size 8 \
+    --num_workers 8 \
     --evaluate_during_training \
     --seed 2048
 
 ```
 
+* 恢复训练
+
+```shell
+export CUDA_VISIBLE_DEVICES=0
+python3 train_re.py \
+    --model_name_or_path "model_path" \
+    --train_data_dir "XFUND/zh_train/image" \
+    --train_label_path "XFUND/zh_train/xfun_normalize_train.json" \
+    --eval_data_dir "XFUND/zh_val/image" \
+    --eval_label_path "XFUND/zh_val/xfun_normalize_val.json" \
+    --label_map_path 'labels/labels_ser.txt' \
+    --num_train_epochs 2 \
+    --eval_steps 10 \
+    --output_dir "output/re/"  \
+    --learning_rate 5e-5 \
+    --warmup_steps 50 \
+    --per_gpu_train_batch_size 8 \
+    --per_gpu_eval_batch_size 8 \
+    --num_workers 8 \
+    --evaluate_during_training \
+    --seed 2048 \
+    --resume
+
+```
+
 最终会打印出`precision`, `recall`, `f1`等指标，模型和训练日志会保存在`./output/re/`文件夹中。
 
+* 评估
+```shell
+export CUDA_VISIBLE_DEVICES=0
+python3 eval_re.py \
+    --model_name_or_path "PP-Layout_v1.0_re_pretrained/" \
+    --max_seq_length 512 \
+    --eval_data_dir "XFUND/zh_val/image" \
+    --eval_label_path "XFUND/zh_val/xfun_normalize_val.json" \
+    --label_map_path 'labels/labels_ser.txt' \
+    --output_dir "output/re/"  \
+    --per_gpu_eval_batch_size 8 \
+    --num_workers 8 \
+    --seed 2048
+```
+最终会打印出`precision`, `recall`, `f1`等指标
+
+
 * 使用评估集合中提供的OCR识别结果进行预测
 
 ```shell
 export CUDA_VISIBLE_DEVICES=0
 python3 infer_re.py \
-    --model_name_or_path "./PP-Layout_v1.0_re_pretrained/" \
+    --model_name_or_path "PP-Layout_v1.0_re_pretrained/" \
     --max_seq_length 512 \
     --eval_data_dir "XFUND/zh_val/image" \
     --eval_label_path "XFUND/zh_val/xfun_normalize_val.json" \
     --label_map_path 'labels/labels_ser.txt' \
-    --output_dir "output_res"  \
+    --output_dir "output/re/"  \
     --per_gpu_eval_batch_size 1 \
     --seed 2048
 ```
@@ -231,11 +313,12 @@ python3 infer_re.py \
 
 ```shell
 export CUDA_VISIBLE_DEVICES=0
-# python3.7 infer_ser_re_e2e.py \
-    --model_name_or_path "./PP-Layout_v1.0_ser_pretrained/" \
-    --re_model_name_or_path "./PP-Layout_v1.0_re_pretrained/" \
+python3.7 infer_ser_re_e2e.py \
+    --model_name_or_path "PP-Layout_v1.0_ser_pretrained/" \
+    --re_model_name_or_path "PP-Layout_v1.0_re_pretrained/" \
+    --ser_model_type "LayoutXLM" \
     --max_seq_length 512 \
-    --output_dir "output_ser_re_e2e_train/" \
+    --output_dir "output/ser_re_e2e/" \
     --infer_imgs "images/input/zh_val_21.jpg"
 ```
 

ppstructure/vqa/eval_re.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(__dir__)
+sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
+
+import paddle
+
+from paddlenlp.transformers import LayoutXLMTokenizer, LayoutXLMModel, LayoutXLMForRelationExtraction
+
+from xfun import XFUNDataset
+from utils import parse_args, get_bio_label_maps, print_arguments
+from data_collator import DataCollator
+from metric import re_score
+
+from ppocr.utils.logging import get_logger
+
+
+def cal_metric(re_preds, re_labels, entities):
+    gt_relations = []
+    for b in range(len(re_labels)):
+        rel_sent = []
+        for head, tail in zip(re_labels[b]["head"], re_labels[b]["tail"]):
+            rel = {}
+            rel["head_id"] = head
+            rel["head"] = (entities[b]["start"][rel["head_id"]],
+                           entities[b]["end"][rel["head_id"]])
+            rel["head_type"] = entities[b]["label"][rel["head_id"]]
+
+            rel["tail_id"] = tail
+            rel["tail"] = (entities[b]["start"][rel["tail_id"]],
+                           entities[b]["end"][rel["tail_id"]])
+            rel["tail_type"] = entities[b]["label"][rel["tail_id"]]
+
+            rel["type"] = 1
+            rel_sent.append(rel)
+        gt_relations.append(rel_sent)
+    re_metrics = re_score(re_preds, gt_relations, mode="boundaries")
+    return re_metrics
+
+
+def evaluate(model, eval_dataloader, logger, prefix=""):
+    # Eval!
+    logger.info("***** Running evaluation {} *****".format(prefix))
+    logger.info("  Num examples = {}".format(len(eval_dataloader.dataset)))
+
+    re_preds = []
+    re_labels = []
+    entities = []
+    eval_loss = 0.0
+    model.eval()
+    for idx, batch in enumerate(eval_dataloader):
+        with paddle.no_grad():
+            outputs = model(**batch)
+            loss = outputs['loss'].mean().item()
+            if paddle.distributed.get_rank() == 0:
+                logger.info("[Eval] process: {}/{}, loss: {:.5f}".format(
+                    idx, len(eval_dataloader), loss))
+
+            eval_loss += loss
+        re_preds.extend(outputs['pred_relations'])
+        re_labels.extend(batch['relations'])
+        entities.extend(batch['entities'])
+    re_metrics = cal_metric(re_preds, re_labels, entities)
+    re_metrics = {
+        "precision": re_metrics["ALL"]["p"],
+        "recall": re_metrics["ALL"]["r"],
+        "f1": re_metrics["ALL"]["f1"],
+    }
+    model.train()
+    return re_metrics
+
+
+def eval(args):
+    logger = get_logger()
+    label2id_map, id2label_map = get_bio_label_maps(args.label_map_path)
+    pad_token_label_id = paddle.nn.CrossEntropyLoss().ignore_index
+
+    tokenizer = LayoutXLMTokenizer.from_pretrained(args.model_name_or_path)
+
+    model = LayoutXLMForRelationExtraction.from_pretrained(
+        args.model_name_or_path)
+
+    eval_dataset = XFUNDataset(
+        tokenizer,
+        data_dir=args.eval_data_dir,
+        label_path=args.eval_label_path,
+        label2id_map=label2id_map,
+        img_size=(224, 224),
+        max_seq_len=args.max_seq_length,
+        pad_token_label_id=pad_token_label_id,
+        contains_re=True,
+        add_special_ids=False,
+        return_attention_mask=True,
+        load_mode='all')
+
+    eval_dataloader = paddle.io.DataLoader(
+        eval_dataset,
+        batch_size=args.per_gpu_eval_batch_size,
+        num_workers=args.num_workers,
+        shuffle=False,
+        collate_fn=DataCollator())
+
+    results = evaluate(model, eval_dataloader, logger)
+    logger.info("eval results: {}".format(results))
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    eval(args)

ppstructure/vqa/eval_ser.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(__dir__)
+sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
+
+import random
+import time
+import copy
+import logging
+
+import argparse
+import paddle
+import numpy as np
+from seqeval.metrics import classification_report, f1_score, precision_score, recall_score
+from paddlenlp.transformers import LayoutXLMModel, LayoutXLMTokenizer, LayoutXLMForTokenClassification
+from paddlenlp.transformers import LayoutLMModel, LayoutLMTokenizer, LayoutLMForTokenClassification
+
+from xfun import XFUNDataset
+from losses import SERLoss
+from utils import parse_args, get_bio_label_maps, print_arguments
+
+from ppocr.utils.logging import get_logger
+
+MODELS = {
+    'LayoutXLM':
+    (LayoutXLMTokenizer, LayoutXLMModel, LayoutXLMForTokenClassification),
+    'LayoutLM':
+    (LayoutLMTokenizer, LayoutLMModel, LayoutLMForTokenClassification)
+}
+
+
+def eval(args):
+    logger = get_logger()
+    print_arguments(args, logger)
+
+    label2id_map, id2label_map = get_bio_label_maps(args.label_map_path)
+    pad_token_label_id = paddle.nn.CrossEntropyLoss().ignore_index
+
+    tokenizer_class, base_model_class, model_class = MODELS[args.ser_model_type]
+    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
+    model = model_class.from_pretrained(args.model_name_or_path)
+
+    eval_dataset = XFUNDataset(
+        tokenizer,
+        data_dir=args.eval_data_dir,
+        label_path=args.eval_label_path,
+        label2id_map=label2id_map,
+        img_size=(224, 224),
+        pad_token_label_id=pad_token_label_id,
+        contains_re=False,
+        add_special_ids=False,
+        return_attention_mask=True,
+        load_mode='all')
+
+    eval_dataloader = paddle.io.DataLoader(
+        eval_dataset,
+        batch_size=args.per_gpu_eval_batch_size,
+        num_workers=args.num_workers,
+        use_shared_memory=True,
+        collate_fn=None, )
+
+    loss_class = SERLoss(len(label2id_map))
+
+    results, _ = evaluate(args, model, tokenizer, loss_class, eval_dataloader,
+                          label2id_map, id2label_map, pad_token_label_id,
+                          logger)
+
+    logger.info(results)
+
+
+def evaluate(args,
+             model,
+             tokenizer,
+             loss_class,
+             eval_dataloader,
+             label2id_map,
+             id2label_map,
+             pad_token_label_id,
+             logger,
+             prefix=""):
+
+    eval_loss = 0.0
+    nb_eval_steps = 0
+    preds = None
+    out_label_ids = None
+    model.eval()
+    for idx, batch in enumerate(eval_dataloader):
+        with paddle.no_grad():
+            if args.ser_model_type == 'LayoutLM':
+                if 'image' in batch:
+                    batch.pop('image')
+            labels = batch.pop('labels')
+            outputs = model(**batch)
+            if args.ser_model_type == 'LayoutXLM':
+                outputs = outputs[0]
+            loss = loss_class(labels, outputs, batch['attention_mask'])
+
+            loss = loss.mean()
+
+            if paddle.distributed.get_rank() == 0:
+                logger.info("[Eval]process: {}/{}, loss: {:.5f}".format(
+                    idx, len(eval_dataloader), loss.numpy()[0]))
+
+            eval_loss += loss.item()
+        nb_eval_steps += 1
+        if preds is None:
+            preds = outputs.numpy()
+            out_label_ids = labels.numpy()
+        else:
+            preds = np.append(preds, outputs.numpy(), axis=0)
+            out_label_ids = np.append(out_label_ids, labels.numpy(), axis=0)
+
+    eval_loss = eval_loss / nb_eval_steps
+    preds = np.argmax(preds, axis=2)
+
+    # label_map = {i: label.upper() for i, label in enumerate(labels)}
+
+    out_label_list = [[] for _ in range(out_label_ids.shape[0])]
+    preds_list = [[] for _ in range(out_label_ids.shape[0])]
+
+    for i in range(out_label_ids.shape[0]):
+        for j in range(out_label_ids.shape[1]):
+            if out_label_ids[i, j] != pad_token_label_id:
+                out_label_list[i].append(id2label_map[out_label_ids[i][j]])
+                preds_list[i].append(id2label_map[preds[i][j]])
+
+    results = {
+        "loss": eval_loss,
+        "precision": precision_score(out_label_list, preds_list),
+        "recall": recall_score(out_label_list, preds_list),
+        "f1": f1_score(out_label_list, preds_list),
+    }
+
+    with open(
+            os.path.join(args.output_dir, "test_gt.txt"), "w",
+            encoding='utf-8') as fout:
+        for lbl in out_label_list:
+            for l in lbl:
+                fout.write(l + "\t")
+            fout.write("\n")
+    with open(
+            os.path.join(args.output_dir, "test_pred.txt"), "w",
+            encoding='utf-8') as fout:
+        for lbl in preds_list:
+            for l in lbl:
+                fout.write(l + "\t")
+            fout.write("\n")
+
+    report = classification_report(out_label_list, preds_list)
+    logger.info("\n" + report)
+
+    logger.info("***** Eval results %s *****", prefix)
+    for key in sorted(results.keys()):
+        logger.info("  %s = %s", key, str(results[key]))
+    model.train()
+    return results, preds_list
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    eval(args)

ppstructure/vqa/helper/eval_with_label_end2end.py

@@ -15,13 +15,12 @@
 import os
 import re
 import sys
-# import Polygon
 import shapely
 from shapely.geometry import Polygon
 import numpy as np
 from collections import defaultdict
 import operator
-import editdistance
+import Levenshtein
 import argparse
 import json
 import copy
@@ -38,7 +37,7 @@ def parse_ser_results_fp(fp, fp_type="gt", ignore_background=True):
     assert fp_type in ["gt", "pred"]
     key = "label" if fp_type == "gt" else "pred"
     res_dict = dict()
-    with open(fp, "r") as fin:
+    with open(fp, "r", encoding='utf-8') as fin:
         lines = fin.readlines()
 
     for _, line in enumerate(lines):
@@ -95,7 +94,7 @@ def ed(args, str1, str2):
     if args.ignore_case:
         str1 = str1.lower()
         str2 = str2.lower()
-    return editdistance.eval(str1, str2)
+    return Levenshtein.distance(str1, str2)
 
 
 def convert_bbox_to_polygon(bbox):
@@ -115,8 +114,6 @@ def eval_e2e(args):
     # pred
     dt_results = parse_ser_results_fp(args.pred_json_path, "pred",
                                       args.ignore_background)
-    assert set(gt_results.keys()) == set(dt_results.keys())
-
     iou_thresh = args.iou_thres
     num_gt_chars = 0
     gt_count = 0
@@ -124,7 +121,7 @@ def eval_e2e(args):
     hit = 0
     ed_sum = 0
 
-    for img_name in gt_results:
+    for img_name in dt_results:
         gt_info = gt_results[img_name]
         gt_count += len(gt_info)
 

ppstructure/vqa/helper/trans_xfun_data.py

@@ -16,13 +16,13 @@ import json
 
 
 def transfer_xfun_data(json_path=None, output_file=None):
-    with open(json_path, "r") as fin:
+    with open(json_path, "r", encoding='utf-8') as fin:
         lines = fin.readlines()
 
     json_info = json.loads(lines[0])
     documents = json_info["documents"]
     label_info = {}
-    with open(output_file, "w") as fout:
+    with open(output_file, "w", encoding='utf-8') as fout:
         for idx, document in enumerate(documents):
             img_info = document["img"]
             document = document["document"]

ppstructure/vqa/infer.sh

+export CUDA_VISIBLE_DEVICES=6
+# python3.7 infer_ser_e2e.py \
+#     --model_name_or_path "output/ser_distributed/best_model" \
+#     --max_seq_length 512 \
+#     --output_dir "output_res_e2e/" \
+#     --infer_imgs "/ssd1/zhoujun20/VQA/data/XFUN_v1.0_data/zh.val/zh_val_0.jpg"
+
+
+# python3.7 infer_ser_re_e2e.py \
+#     --model_name_or_path "output/ser_distributed/best_model" \
+#     --re_model_name_or_path "output/re_test/best_model" \
+#     --max_seq_length 512 \
+#     --output_dir "output_ser_re_e2e_train/" \
+#     --infer_imgs "images/input/zh_val_21.jpg"
+
+# python3.7 infer_ser.py \
+#     --model_name_or_path "output/ser_LayoutLM/best_model" \
+#     --ser_model_type "LayoutLM" \
+#     --output_dir "ser_LayoutLM/" \
+#     --infer_imgs "images/input/zh_val_21.jpg" \
+#     --ocr_json_path "/ssd1/zhoujun20/VQA/data/XFUN_v1.0_data/xfun_normalize_val.json"
+
+python3.7 infer_ser.py \
+    --model_name_or_path "output/ser_new/best_model" \
+    --ser_model_type "LayoutXLM" \
+    --output_dir "ser_new/" \
+    --infer_imgs "images/input/zh_val_21.jpg" \
+    --ocr_json_path "/ssd1/zhoujun20/VQA/data/XFUN_v1.0_data/xfun_normalize_val.json"
+
+# python3.7 infer_ser_e2e.py \
+#     --model_name_or_path "output/ser_new/best_model" \
+#     --ser_model_type "LayoutXLM" \
+#     --max_seq_length 512 \
+#     --output_dir "output/ser_new/" \
+#     --infer_imgs "images/input/zh_val_0.jpg"
+
+
+# python3.7 infer_ser_e2e.py \
+#     --model_name_or_path "output/ser_LayoutLM/best_model" \
+#     --ser_model_type "LayoutLM" \
+#     --max_seq_length 512 \
+#     --output_dir "output/ser_LayoutLM/" \
+#     --infer_imgs "images/input/zh_val_0.jpg"
+
+# python3 infer_re.py \
+#     --model_name_or_path "/ssd1/zhoujun20/VQA/PaddleOCR/ppstructure/vqa/output/re_test/best_model/" \
+#     --max_seq_length 512 \
+#     --eval_data_dir "/ssd1/zhoujun20/VQA/data/XFUN_v1.0_data/zh.val" \
+#     --eval_label_path "/ssd1/zhoujun20/VQA/data/XFUN_v1.0_data/xfun_normalize_val.json" \
+#     --label_map_path 'labels/labels_ser.txt' \
+#     --output_dir "output_res"  \
+#     --per_gpu_eval_batch_size 1 \
+#     --seed 2048
+
+# python3.7 infer_ser_re_e2e.py \
+#     --model_name_or_path "output/ser_LayoutLM/best_model" \
+#     --ser_model_type "LayoutLM" \
+#     --re_model_name_or_path "output/re_new/best_model" \
+#     --max_seq_length 512 \
+#     --output_dir "output_ser_re_e2e/" \
+#     --infer_imgs "images/input/zh_val_21.jpg"
\ No newline at end of file

ppstructure/vqa/infer_re.py

@@ -56,15 +56,19 @@ def infer(args):
     ocr_info_list = load_ocr(args.eval_data_dir, args.eval_label_path)
 
     for idx, batch in enumerate(eval_dataloader):
-        logger.info("[Infer] process: {}/{}".format(idx, len(eval_dataloader)))
-        with paddle.no_grad():
-            outputs = model(**batch)
-        pred_relations = outputs['pred_relations']
-
         ocr_info = ocr_info_list[idx]
         image_path = ocr_info['image_path']
         ocr_info = ocr_info['ocr_info']
 
+        save_img_path = os.path.join(
+            args.output_dir,
+            os.path.splitext(os.path.basename(image_path))[0] + "_re.jpg")
+        logger.info("[Infer] process: {}/{}, save result to {}".format(
+            idx, len(eval_dataloader), save_img_path))
+        with paddle.no_grad():
+            outputs = model(**batch)
+        pred_relations = outputs['pred_relations']
+
         # 根据entity里的信息，做token解码后去过滤不要的ocr_info
         ocr_info = filter_bg_by_txt(ocr_info, batch, tokenizer)
 
@@ -85,14 +89,13 @@ def infer(args):
 
         img = cv2.imread(image_path)
         img_show = draw_re_results(img, result)
-        save_path = os.path.join(args.output_dir, os.path.basename(image_path))
-        cv2.imwrite(save_path, img_show)
+        cv2.imwrite(save_img_path, img_show)
 
 
 def load_ocr(img_folder, json_path):
     import json
     d = []
-    with open(json_path, "r") as fin:
+    with open(json_path, "r", encoding='utf-8') as fin:
         lines = fin.readlines()
         for line in lines:
             image_name, info_str = line.split("\t")

ppstructure/vqa/infer_ser.py

@@ -24,6 +24,14 @@ import paddle
 # relative reference
 from utils import parse_args, get_image_file_list, draw_ser_results, get_bio_label_maps
 from paddlenlp.transformers import LayoutXLMModel, LayoutXLMTokenizer, LayoutXLMForTokenClassification
+from paddlenlp.transformers import LayoutLMModel, LayoutLMTokenizer, LayoutLMForTokenClassification
+
+MODELS = {
+    'LayoutXLM':
+    (LayoutXLMTokenizer, LayoutXLMModel, LayoutXLMForTokenClassification),
+    'LayoutLM':
+    (LayoutLMTokenizer, LayoutLMModel, LayoutLMForTokenClassification)
+}
 
 
 def pad_sentences(tokenizer,
@@ -59,7 +67,8 @@ def pad_sentences(tokenizer,
             encoded_inputs["bbox"] = encoded_inputs["bbox"] + [[0, 0, 0, 0]
                                                                ] * difference
         else:
-            assert False, f"padding_side of tokenizer just supports [\"right\"] but got {tokenizer.padding_side}"
+            assert False, "padding_side of tokenizer just supports [\"right\"] but got {}".format(
+                tokenizer.padding_side)
     else:
         if return_attention_mask:
             encoded_inputs["attention_mask"] = [1] * len(encoded_inputs[
@@ -216,15 +225,15 @@ def infer(args):
     os.makedirs(args.output_dir, exist_ok=True)
 
     # init token and model
-    tokenizer = LayoutXLMTokenizer.from_pretrained(args.model_name_or_path)
-    # model = LayoutXLMModel.from_pretrained(args.model_name_or_path)
-    model = LayoutXLMForTokenClassification.from_pretrained(
-        args.model_name_or_path)
+    tokenizer_class, base_model_class, model_class = MODELS[args.ser_model_type]
+    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
+    model = model_class.from_pretrained(args.model_name_or_path)
+
     model.eval()
 
     # load ocr results json
     ocr_results = dict()
-    with open(args.ocr_json_path, "r") as fin:
+    with open(args.ocr_json_path, "r", encoding='utf-8') as fin:
         lines = fin.readlines()
         for line in lines:
             img_name, json_info = line.split("\t")
@@ -234,9 +243,15 @@ def infer(args):
     infer_imgs = get_image_file_list(args.infer_imgs)
 
     # loop for infer
-    with open(os.path.join(args.output_dir, "infer_results.txt"), "w") as fout:
+    with open(
+            os.path.join(args.output_dir, "infer_results.txt"),
+            "w",
+            encoding='utf-8') as fout:
         for idx, img_path in enumerate(infer_imgs):
-            print("process: [{}/{}]".format(idx, len(infer_imgs), img_path))
+            save_img_path = os.path.join(args.output_dir,
+                                         os.path.basename(img_path))
+            print("process: [{}/{}], save result to {}".format(
+                idx, len(infer_imgs), save_img_path))
 
             img = cv2.imread(img_path)
 
@@ -246,15 +261,21 @@ def infer(args):
                 ori_img=img,
                 ocr_info=ocr_info,
                 max_seq_len=args.max_seq_length)
-
-            outputs = model(
+            if args.ser_model_type == 'LayoutLM':
+                preds = model(
+                    input_ids=inputs["input_ids"],
+                    bbox=inputs["bbox"],
+                    token_type_ids=inputs["token_type_ids"],
+                    attention_mask=inputs["attention_mask"])
+            elif args.ser_model_type == 'LayoutXLM':
+                preds = model(
                     input_ids=inputs["input_ids"],
                     bbox=inputs["bbox"],
                     image=inputs["image"],
                     token_type_ids=inputs["token_type_ids"],
                     attention_mask=inputs["attention_mask"])
+                preds = preds[0]
 
-            preds = outputs[0]
             preds = postprocess(inputs["attention_mask"], preds,
                                 args.label_map_path)
             ocr_info = merge_preds_list_with_ocr_info(
@@ -267,9 +288,7 @@ def infer(args):
                 }, ensure_ascii=False) + "\n")
 
             img_res = draw_ser_results(img, ocr_info)
-            cv2.imwrite(
-                os.path.join(args.output_dir, os.path.basename(img_path)),
-                img_res)
+            cv2.imwrite(save_img_path, img_res)
 
     return
 

ppstructure/vqa/infer_ser_e2e.py

@@ -22,14 +22,20 @@ from PIL import Image
 
 import paddle
 from paddlenlp.transformers import LayoutXLMModel, LayoutXLMTokenizer, LayoutXLMForTokenClassification
-
-from paddleocr import PaddleOCR
+from paddlenlp.transformers import LayoutLMModel, LayoutLMTokenizer, LayoutLMForTokenClassification
 
 # relative reference
 from utils import parse_args, get_image_file_list, draw_ser_results, get_bio_label_maps
 
 from utils import pad_sentences, split_page, preprocess, postprocess, merge_preds_list_with_ocr_info
 
+MODELS = {
+    'LayoutXLM':
+    (LayoutXLMTokenizer, LayoutXLMModel, LayoutXLMForTokenClassification),
+    'LayoutLM':
+    (LayoutLMTokenizer, LayoutLMModel, LayoutLMForTokenClassification)
+}
+
 
 def trans_poly_to_bbox(poly):
     x1 = np.min([p[0] for p in poly])
@@ -52,19 +58,23 @@ def parse_ocr_info_for_ser(ocr_result):
 
 class SerPredictor(object):
     def __init__(self, args):
+        self.args = args
         self.max_seq_length = args.max_seq_length
 
         # init ser token and model
-        self.tokenizer = LayoutXLMTokenizer.from_pretrained(
-            args.model_name_or_path)
-        self.model = LayoutXLMForTokenClassification.from_pretrained(
+        tokenizer_class, base_model_class, model_class = MODELS[
+            args.ser_model_type]
+        self.tokenizer = tokenizer_class.from_pretrained(
             args.model_name_or_path)
+        self.model = model_class.from_pretrained(args.model_name_or_path)
         self.model.eval()
 
         # init ocr_engine
+        from paddleocr import PaddleOCR
+
         self.ocr_engine = PaddleOCR(
-            rec_model_dir=args.ocr_rec_model_dir,
-            det_model_dir=args.ocr_det_model_dir,
+            rec_model_dir=args.rec_model_dir,
+            det_model_dir=args.det_model_dir,
             use_angle_cls=False,
             show_log=False)
         # init dict
@@ -88,14 +98,21 @@ class SerPredictor(object):
             ocr_info=ocr_info,
             max_seq_len=self.max_seq_length)
 
-        outputs = self.model(
+        if self.args.ser_model_type == 'LayoutLM':
+            preds = self.model(
+                input_ids=inputs["input_ids"],
+                bbox=inputs["bbox"],
+                token_type_ids=inputs["token_type_ids"],
+                attention_mask=inputs["attention_mask"])
+        elif self.args.ser_model_type == 'LayoutXLM':
+            preds = self.model(
                 input_ids=inputs["input_ids"],
                 bbox=inputs["bbox"],
                 image=inputs["image"],
                 token_type_ids=inputs["token_type_ids"],
                 attention_mask=inputs["attention_mask"])
+            preds = preds[0]
 
-        preds = outputs[0]
         preds = postprocess(inputs["attention_mask"], preds, self.id2label_map)
         ocr_info = merge_preds_list_with_ocr_info(
             ocr_info, inputs["segment_offset_id"], preds,
@@ -112,9 +129,16 @@ if __name__ == "__main__":
 
     # loop for infer
     ser_engine = SerPredictor(args)
-    with open(os.path.join(args.output_dir, "infer_results.txt"), "w") as fout:
+    with open(
+            os.path.join(args.output_dir, "infer_results.txt"),
+            "w",
+            encoding='utf-8') as fout:
         for idx, img_path in enumerate(infer_imgs):
-            print("process: [{}/{}], {}".format(idx, len(infer_imgs), img_path))
+            save_img_path = os.path.join(
+                args.output_dir,
+                os.path.splitext(os.path.basename(img_path))[0] + "_ser.jpg")
+            print("process: [{}/{}], save result to {}".format(
+                idx, len(infer_imgs), save_img_path))
 
             img = cv2.imread(img_path)
 
@@ -125,7 +149,4 @@ if __name__ == "__main__":
                 }, ensure_ascii=False) + "\n")
 
             img_res = draw_ser_results(img, result)
-            cv2.imwrite(
-                os.path.join(args.output_dir,
-                             os.path.splitext(os.path.basename(img_path))[0] +
-                             "_ser.jpg"), img_res)
+            cv2.imwrite(save_img_path, img_res)

ppstructure/vqa/infer_ser_re_e2e.py

@@ -112,9 +112,16 @@ if __name__ == "__main__":
 
     # loop for infer
     ser_re_engine = SerReSystem(args)
-    with open(os.path.join(args.output_dir, "infer_results.txt"), "w") as fout:
+    with open(
+            os.path.join(args.output_dir, "infer_results.txt"),
+            "w",
+            encoding='utf-8') as fout:
         for idx, img_path in enumerate(infer_imgs):
-            print("process: [{}/{}], {}".format(idx, len(infer_imgs), img_path))
+            save_img_path = os.path.join(
+                args.output_dir,
+                os.path.splitext(os.path.basename(img_path))[0] + "_re.jpg")
+            print("process: [{}/{}], save result to {}".format(
+                idx, len(infer_imgs), save_img_path))
 
             img = cv2.imread(img_path)
 
@@ -125,7 +132,4 @@ if __name__ == "__main__":
                 }, ensure_ascii=False) + "\n")
 
             img_res = draw_re_results(img, result)
-            cv2.imwrite(
-                os.path.join(args.output_dir,
-                             os.path.splitext(os.path.basename(img_path))[0] +
-                             "_re.jpg"), img_res)
+            cv2.imwrite(save_img_path, img_res)

ppstructure/vqa/losses.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle import nn
+
+
+class SERLoss(nn.Layer):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.loss_class = nn.CrossEntropyLoss()
+        self.num_classes = num_classes
+        self.ignore_index = self.loss_class.ignore_index
+
+    def forward(self, labels, outputs, attention_mask):
+        if attention_mask is not None:
+            active_loss = attention_mask.reshape([-1, ]) == 1
+            active_outputs = outputs.reshape(
+                [-1, self.num_classes])[active_loss]
+            active_labels = labels.reshape([-1, ])[active_loss]
+            loss = self.loss_class(active_outputs, active_labels)
+        else:
+            loss = self.loss_class(
+                outputs.reshape([-1, self.num_classes]), labels.reshape([-1, ]))
+        return loss

ppstructure/vqa/requirements.txt

 sentencepiece
 yacs
+seqeval
\ No newline at end of file

ppstructure/vqa/train_re.py

@@ -20,82 +20,25 @@ sys.path.append(__dir__)
 sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
 
 import random
+import time
 import numpy as np
 import paddle
 
 from paddlenlp.transformers import LayoutXLMTokenizer, LayoutXLMModel, LayoutXLMForRelationExtraction
 
 from xfun import XFUNDataset
-from utils import parse_args, get_bio_label_maps, print_arguments
+from utils import parse_args, get_bio_label_maps, print_arguments, set_seed
 from data_collator import DataCollator
-from metric import re_score
+from eval_re import evaluate
 
 from ppocr.utils.logging import get_logger
 
 
-def set_seed(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    paddle.seed(seed)
-
-
-def cal_metric(re_preds, re_labels, entities):
-    gt_relations = []
-    for b in range(len(re_labels)):
-        rel_sent = []
-        for head, tail in zip(re_labels[b]["head"], re_labels[b]["tail"]):
-            rel = {}
-            rel["head_id"] = head
-            rel["head"] = (entities[b]["start"][rel["head_id"]],
-                           entities[b]["end"][rel["head_id"]])
-            rel["head_type"] = entities[b]["label"][rel["head_id"]]
-
-            rel["tail_id"] = tail
-            rel["tail"] = (entities[b]["start"][rel["tail_id"]],
-                           entities[b]["end"][rel["tail_id"]])
-            rel["tail_type"] = entities[b]["label"][rel["tail_id"]]
-
-            rel["type"] = 1
-            rel_sent.append(rel)
-        gt_relations.append(rel_sent)
-    re_metrics = re_score(re_preds, gt_relations, mode="boundaries")
-    return re_metrics
-
-
-def evaluate(model, eval_dataloader, logger, prefix=""):
-    # Eval!
-    logger.info("***** Running evaluation {} *****".format(prefix))
-    logger.info("  Num examples = {}".format(len(eval_dataloader.dataset)))
-
-    re_preds = []
-    re_labels = []
-    entities = []
-    eval_loss = 0.0
-    model.eval()
-    for idx, batch in enumerate(eval_dataloader):
-        with paddle.no_grad():
-            outputs = model(**batch)
-            loss = outputs['loss'].mean().item()
-            if paddle.distributed.get_rank() == 0:
-                logger.info("[Eval] process: {}/{}, loss: {:.5f}".format(
-                    idx, len(eval_dataloader), loss))
-
-            eval_loss += loss
-        re_preds.extend(outputs['pred_relations'])
-        re_labels.extend(batch['relations'])
-        entities.extend(batch['entities'])
-    re_metrics = cal_metric(re_preds, re_labels, entities)
-    re_metrics = {
-        "precision": re_metrics["ALL"]["p"],
-        "recall": re_metrics["ALL"]["r"],
-        "f1": re_metrics["ALL"]["f1"],
-    }
-    model.train()
-    return re_metrics
-
-
 def train(args):
     logger = get_logger(log_file=os.path.join(args.output_dir, "train.log"))
+    rank = paddle.distributed.get_rank()
+    distributed = paddle.distributed.get_world_size() > 1
+
     print_arguments(args, logger)
 
     # Added here for reproducibility (even between python 2 and 3)
@@ -105,17 +48,22 @@ def train(args):
     pad_token_label_id = paddle.nn.CrossEntropyLoss().ignore_index
 
     # dist mode
-    if paddle.distributed.get_world_size() > 1:
+    if distributed:
         paddle.distributed.init_parallel_env()
 
     tokenizer = LayoutXLMTokenizer.from_pretrained(args.model_name_or_path)
-
+    if not args.resume:
         model = LayoutXLMModel.from_pretrained(args.model_name_or_path)
         model = LayoutXLMForRelationExtraction(model, dropout=None)
+        logger.info('train from scratch')
+    else:
+        logger.info('resume from {}'.format(args.model_name_or_path))
+        model = LayoutXLMForRelationExtraction.from_pretrained(
+            args.model_name_or_path)
 
     # dist mode
-    if paddle.distributed.get_world_size() > 1:
-        model = paddle.distributed.DataParallel(model)
+    if distributed:
+        model = paddle.DataParallel(model)
 
     train_dataset = XFUNDataset(
         tokenizer,
@@ -145,19 +93,18 @@ def train(args):
 
     train_sampler = paddle.io.DistributedBatchSampler(
         train_dataset, batch_size=args.per_gpu_train_batch_size, shuffle=True)
-    args.train_batch_size = args.per_gpu_train_batch_size * \
-                            max(1, paddle.distributed.get_world_size())
+
     train_dataloader = paddle.io.DataLoader(
         train_dataset,
         batch_sampler=train_sampler,
-        num_workers=8,
+        num_workers=args.num_workers,
         use_shared_memory=True,
         collate_fn=DataCollator())
 
     eval_dataloader = paddle.io.DataLoader(
         eval_dataset,
         batch_size=args.per_gpu_eval_batch_size,
-        num_workers=8,
+        num_workers=args.num_workers,
         shuffle=False,
         collate_fn=DataCollator())
 
@@ -191,7 +138,8 @@ def train(args):
         args.per_gpu_train_batch_size))
     logger.info(
         "  Total train batch size (w. parallel, distributed & accumulation) = {}".
-        format(args.train_batch_size * paddle.distributed.get_world_size()))
+        format(args.per_gpu_train_batch_size *
+               paddle.distributed.get_world_size()))
     logger.info("  Total optimization steps = {}".format(t_total))
 
     global_step = 0
@@ -200,58 +148,78 @@ def train(args):
     best_metirc = {'f1': 0}
     model.train()
 
+    train_reader_cost = 0.0
+    train_run_cost = 0.0
+    total_samples = 0
+    reader_start = time.time()
+
+    print_step = 1
+
     for epoch in range(int(args.num_train_epochs)):
         for step, batch in enumerate(train_dataloader):
+            train_reader_cost += time.time() - reader_start
+            train_start = time.time()
             outputs = model(**batch)
+            train_run_cost += time.time() - train_start
             # model outputs are always tuple in ppnlp (see doc)
             loss = outputs['loss']
             loss = loss.mean()
 
-            logger.info(
-                "epoch: [{}/{}], iter: [{}/{}], global_step:{}, train loss: {}, lr: {}".
-                format(epoch, args.num_train_epochs, step, train_dataloader_len,
-                       global_step, np.mean(loss.numpy()), optimizer.get_lr()))
-
             loss.backward()
             optimizer.step()
             optimizer.clear_grad()
             # lr_scheduler.step()  # Update learning rate schedule
 
             global_step += 1
+            total_samples += batch['image'].shape[0]
 
-            if (paddle.distributed.get_rank() == 0 and args.eval_steps > 0 and
-                    global_step % args.eval_steps == 0):
+            if rank == 0 and step % print_step == 0:
+                logger.info(
+                    "epoch: [{}/{}], iter: [{}/{}], global_step:{}, train loss: {:.6f}, lr: {:.6f}, avg_reader_cost: {:.5f} sec, avg_batch_cost: {:.5f} sec, avg_samples: {:.5f}, ips: {:.5f} images/sec".
+                    format(epoch, args.num_train_epochs, step,
+                           train_dataloader_len, global_step,
+                           np.mean(loss.numpy()),
+                           optimizer.get_lr(), train_reader_cost / print_step, (
+                               train_reader_cost + train_run_cost) / print_step,
+                           total_samples / print_step, total_samples / (
+                               train_reader_cost + train_run_cost)))
+
+                train_reader_cost = 0.0
+                train_run_cost = 0.0
+                total_samples = 0
+
+            if rank == 0 and args.eval_steps > 0 and global_step % args.eval_steps == 0 and args.evaluate_during_training:
                 # Log metrics
-                if (paddle.distributed.get_rank() == 0 and args.
-                        evaluate_during_training):  # Only evaluate when single GPU otherwise metrics may not average well
+                # Only evaluate when single GPU otherwise metrics may not average well
                 results = evaluate(model, eval_dataloader, logger)
-                    if results['f1'] > best_metirc['f1']:
+                if results['f1'] >= best_metirc['f1']:
                     best_metirc = results
-                        output_dir = os.path.join(args.output_dir,
-                                                  "checkpoint-best")
+                    output_dir = os.path.join(args.output_dir, "best_model")
                     os.makedirs(output_dir, exist_ok=True)
+                    if distributed:
+                        model._layers.save_pretrained(output_dir)
+                    else:
                         model.save_pretrained(output_dir)
                     tokenizer.save_pretrained(output_dir)
                     paddle.save(args,
-                                    os.path.join(output_dir,
-                                                 "training_args.bin"))
+                                os.path.join(output_dir, "training_args.bin"))
                     logger.info("Saving model checkpoint to {}".format(
                         output_dir))
                 logger.info("eval results: {}".format(results))
                 logger.info("best_metirc: {}".format(best_metirc))
+            reader_start = time.time()
 
-            if (paddle.distributed.get_rank() == 0 and args.save_steps > 0 and
-                    global_step % args.save_steps == 0):
+        if rank == 0:
             # Save model checkpoint
-                output_dir = os.path.join(args.output_dir, "checkpoint-latest")
+            output_dir = os.path.join(args.output_dir, "latest_model")
             os.makedirs(output_dir, exist_ok=True)
-                if paddle.distributed.get_rank() == 0:
+            if distributed:
+                model._layers.save_pretrained(output_dir)
+            else:
                 model.save_pretrained(output_dir)
             tokenizer.save_pretrained(output_dir)
-                    paddle.save(args,
-                                os.path.join(output_dir, "training_args.bin"))
-                    logger.info("Saving model checkpoint to {}".format(
-                        output_dir))
+            paddle.save(args, os.path.join(output_dir, "training_args.bin"))
+            logger.info("Saving model checkpoint to {}".format(output_dir))
     logger.info("best_metirc: {}".format(best_metirc))
 
 

ppstructure/vqa/train_ser.py

@@ -20,6 +20,7 @@ sys.path.append(__dir__)
 sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
 
 import random
+import time
 import copy
 import logging
 
@@ -28,39 +29,52 @@ import paddle
 import numpy as np
 from seqeval.metrics import classification_report, f1_score, precision_score, recall_score
 from paddlenlp.transformers import LayoutXLMModel, LayoutXLMTokenizer, LayoutXLMForTokenClassification
-from xfun import XFUNDataset
-from utils import parse_args
-from utils import get_bio_label_maps
-from utils import print_arguments
+from paddlenlp.transformers import LayoutLMModel, LayoutLMTokenizer, LayoutLMForTokenClassification
 
+from xfun import XFUNDataset
+from utils import parse_args, get_bio_label_maps, print_arguments, set_seed
+from eval_ser import evaluate
+from losses import SERLoss
 from ppocr.utils.logging import get_logger
 
-
-def set_seed(args):
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    paddle.seed(args.seed)
+MODELS = {
+    'LayoutXLM':
+    (LayoutXLMTokenizer, LayoutXLMModel, LayoutXLMForTokenClassification),
+    'LayoutLM':
+    (LayoutLMTokenizer, LayoutLMModel, LayoutLMForTokenClassification)
+}
 
 
 def train(args):
     os.makedirs(args.output_dir, exist_ok=True)
+    rank = paddle.distributed.get_rank()
+    distributed = paddle.distributed.get_world_size() > 1
+
     logger = get_logger(log_file=os.path.join(args.output_dir, "train.log"))
     print_arguments(args, logger)
 
     label2id_map, id2label_map = get_bio_label_maps(args.label_map_path)
-    pad_token_label_id = paddle.nn.CrossEntropyLoss().ignore_index
+    loss_class = SERLoss(len(label2id_map))
+
+    pad_token_label_id = loss_class.ignore_index
 
     # dist mode
-    if paddle.distributed.get_world_size() > 1:
+    if distributed:
         paddle.distributed.init_parallel_env()
 
-    tokenizer = LayoutXLMTokenizer.from_pretrained(args.model_name_or_path)
-    base_model = LayoutXLMModel.from_pretrained(args.model_name_or_path)
-    model = LayoutXLMForTokenClassification(
+    tokenizer_class, base_model_class, model_class = MODELS[args.ser_model_type]
+    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
+    if not args.resume:
+        base_model = base_model_class.from_pretrained(args.model_name_or_path)
+        model = model_class(
             base_model, num_classes=len(label2id_map), dropout=None)
+        logger.info('train from scratch')
+    else:
+        logger.info('resume from {}'.format(args.model_name_or_path))
+        model = model_class.from_pretrained(args.model_name_or_path)
 
     # dist mode
-    if paddle.distributed.get_world_size() > 1:
+    if distributed:
         model = paddle.DataParallel(model)
 
     train_dataset = XFUNDataset(
@@ -74,17 +88,32 @@ def train(args):
         add_special_ids=False,
         return_attention_mask=True,
         load_mode='all')
+    eval_dataset = XFUNDataset(
+        tokenizer,
+        data_dir=args.eval_data_dir,
+        label_path=args.eval_label_path,
+        label2id_map=label2id_map,
+        img_size=(224, 224),
+        pad_token_label_id=pad_token_label_id,
+        contains_re=False,
+        add_special_ids=False,
+        return_attention_mask=True,
+        load_mode='all')
 
     train_sampler = paddle.io.DistributedBatchSampler(
         train_dataset, batch_size=args.per_gpu_train_batch_size, shuffle=True)
 
-    args.train_batch_size = args.per_gpu_train_batch_size * max(
-        1, paddle.distributed.get_world_size())
-
     train_dataloader = paddle.io.DataLoader(
         train_dataset,
         batch_sampler=train_sampler,
-        num_workers=0,
+        num_workers=args.num_workers,
+        use_shared_memory=True,
+        collate_fn=None, )
+
+    eval_dataloader = paddle.io.DataLoader(
+        eval_dataset,
+        batch_size=args.per_gpu_eval_batch_size,
+        num_workers=args.num_workers,
         use_shared_memory=True,
         collate_fn=None, )
 
@@ -117,182 +146,103 @@ def train(args):
                 args.per_gpu_train_batch_size)
     logger.info(
         "  Total train batch size (w. parallel, distributed) = %d",
-        args.train_batch_size * paddle.distributed.get_world_size(), )
+        args.per_gpu_train_batch_size * paddle.distributed.get_world_size(), )
     logger.info("  Total optimization steps = %d", t_total)
 
     global_step = 0
     tr_loss = 0.0
-    set_seed(args)
+    set_seed(args.seed)
     best_metrics = None
 
+    train_reader_cost = 0.0
+    train_run_cost = 0.0
+    total_samples = 0
+    reader_start = time.time()
+
+    print_step = 1
+    model.train()
     for epoch_id in range(args.num_train_epochs):
         for step, batch in enumerate(train_dataloader):
-            model.train()
+            train_reader_cost += time.time() - reader_start
+
+            if args.ser_model_type == 'LayoutLM':
+                if 'image' in batch:
+                    batch.pop('image')
+            labels = batch.pop('labels')
+
+            train_start = time.time()
             outputs = model(**batch)
+            train_run_cost += time.time() - train_start
+            if args.ser_model_type == 'LayoutXLM':
+                outputs = outputs[0]
+            loss = loss_class(labels, outputs, batch['attention_mask'])
+
             # model outputs are always tuple in ppnlp (see doc)
-            loss = outputs[0]
             loss = loss.mean()
-            logger.info(
-                "epoch: [{}/{}], iter: [{}/{}], global_step:{}, train loss: {}, lr: {}".
-                format(epoch_id, args.num_train_epochs, step,
-                       len(train_dataloader), global_step,
-                       loss.numpy()[0], lr_scheduler.get_lr()))
-
             loss.backward()
             tr_loss += loss.item()
             optimizer.step()
             lr_scheduler.step()  # Update learning rate schedule
             optimizer.clear_grad()
             global_step += 1
+            total_samples += batch['input_ids'].shape[0]
+
+            if rank == 0 and step % print_step == 0:
+                logger.info(
+                    "epoch: [{}/{}], iter: [{}/{}], global_step:{}, train loss: {:.6f}, lr: {:.6f}, avg_reader_cost: {:.5f} sec, avg_batch_cost: {:.5f} sec, avg_samples: {:.5f}, ips: {:.5f} images/sec".
+                    format(epoch_id, args.num_train_epochs, step,
+                           len(train_dataloader), global_step,
+                           loss.numpy()[0],
+                           lr_scheduler.get_lr(), train_reader_cost /
+                           print_step, (train_reader_cost + train_run_cost) /
+                           print_step, total_samples / print_step, total_samples
+                           / (train_reader_cost + train_run_cost)))
+
+                train_reader_cost = 0.0
+                train_run_cost = 0.0
+                total_samples = 0
 
-            if (paddle.distributed.get_rank() == 0 and args.eval_steps > 0 and
-                    global_step % args.eval_steps == 0):
+            if rank == 0 and args.eval_steps > 0 and global_step % args.eval_steps == 0 and args.evaluate_during_training:
                 # Log metrics
                 # Only evaluate when single GPU otherwise metrics may not average well
-                if paddle.distributed.get_rank(
-                ) == 0 and args.evaluate_during_training:
-                    results, _ = evaluate(args, model, tokenizer, label2id_map,
-                                          id2label_map, pad_token_label_id,
-                                          logger)
-
-                    if best_metrics is None or results["f1"] >= best_metrics[
-                            "f1"]:
+                results, _ = evaluate(args, model, tokenizer, loss_class,
+                                      eval_dataloader, label2id_map,
+                                      id2label_map, pad_token_label_id, logger)
+
+                if best_metrics is None or results["f1"] >= best_metrics["f1"]:
                     best_metrics = copy.deepcopy(results)
                     output_dir = os.path.join(args.output_dir, "best_model")
                     os.makedirs(output_dir, exist_ok=True)
-                        if paddle.distributed.get_rank() == 0:
+                    if distributed:
+                        model._layers.save_pretrained(output_dir)
+                    else:
                         model.save_pretrained(output_dir)
                     tokenizer.save_pretrained(output_dir)
-                            paddle.save(
-                                args,
+                    paddle.save(args,
                                 os.path.join(output_dir, "training_args.bin"))
-                            logger.info("Saving model checkpoint to %s",
-                                        output_dir)
+                    logger.info("Saving model checkpoint to {}".format(
+                        output_dir))
 
                 logger.info("[epoch {}/{}][iter: {}/{}] results: {}".format(
                     epoch_id, args.num_train_epochs, step,
                     len(train_dataloader), results))
                 if best_metrics is not None:
                     logger.info("best metrics: {}".format(best_metrics))
-
-            if paddle.distributed.get_rank(
-            ) == 0 and args.save_steps > 0 and global_step % args.save_steps == 0:
+            reader_start = time.time()
+        if rank == 0:
             # Save model checkpoint
-                output_dir = os.path.join(args.output_dir,
-                                          "checkpoint-{}".format(global_step))
+            output_dir = os.path.join(args.output_dir, "latest_model")
             os.makedirs(output_dir, exist_ok=True)
-                if paddle.distributed.get_rank() == 0:
+            if distributed:
+                model._layers.save_pretrained(output_dir)
+            else:
                 model.save_pretrained(output_dir)
             tokenizer.save_pretrained(output_dir)
-                    paddle.save(args,
-                                os.path.join(output_dir, "training_args.bin"))
-                    logger.info("Saving model checkpoint to %s", output_dir)
-
+            paddle.save(args, os.path.join(output_dir, "training_args.bin"))
+            logger.info("Saving model checkpoint to {}".format(output_dir))
     return global_step, tr_loss / global_step
 
 
-def evaluate(args,
-             model,
-             tokenizer,
-             label2id_map,
-             id2label_map,
-             pad_token_label_id,
-             logger,
-             prefix=""):
-    eval_dataset = XFUNDataset(
-        tokenizer,
-        data_dir=args.eval_data_dir,
-        label_path=args.eval_label_path,
-        label2id_map=label2id_map,
-        img_size=(224, 224),
-        pad_token_label_id=pad_token_label_id,
-        contains_re=False,
-        add_special_ids=False,
-        return_attention_mask=True,
-        load_mode='all')
-
-    args.eval_batch_size = args.per_gpu_eval_batch_size * max(
-        1, paddle.distributed.get_world_size())
-
-    eval_dataloader = paddle.io.DataLoader(
-        eval_dataset,
-        batch_size=args.eval_batch_size,
-        num_workers=0,
-        use_shared_memory=True,
-        collate_fn=None, )
-
-    # Eval!
-    logger.info("***** Running evaluation %s *****", prefix)
-    logger.info("  Num examples = %d", len(eval_dataset))
-    logger.info("  Batch size = %d", args.eval_batch_size)
-    eval_loss = 0.0
-    nb_eval_steps = 0
-    preds = None
-    out_label_ids = None
-    model.eval()
-    for idx, batch in enumerate(eval_dataloader):
-        with paddle.no_grad():
-            outputs = model(**batch)
-            tmp_eval_loss, logits = outputs[:2]
-
-            tmp_eval_loss = tmp_eval_loss.mean()
-
-            if paddle.distributed.get_rank() == 0:
-                logger.info("[Eval]process: {}/{}, loss: {:.5f}".format(
-                    idx, len(eval_dataloader), tmp_eval_loss.numpy()[0]))
-
-            eval_loss += tmp_eval_loss.item()
-        nb_eval_steps += 1
-        if preds is None:
-            preds = logits.numpy()
-            out_label_ids = batch["labels"].numpy()
-        else:
-            preds = np.append(preds, logits.numpy(), axis=0)
-            out_label_ids = np.append(
-                out_label_ids, batch["labels"].numpy(), axis=0)
-
-    eval_loss = eval_loss / nb_eval_steps
-    preds = np.argmax(preds, axis=2)
-
-    # label_map = {i: label.upper() for i, label in enumerate(labels)}
-
-    out_label_list = [[] for _ in range(out_label_ids.shape[0])]
-    preds_list = [[] for _ in range(out_label_ids.shape[0])]
-
-    for i in range(out_label_ids.shape[0]):
-        for j in range(out_label_ids.shape[1]):
-            if out_label_ids[i, j] != pad_token_label_id:
-                out_label_list[i].append(id2label_map[out_label_ids[i][j]])
-                preds_list[i].append(id2label_map[preds[i][j]])
-
-    results = {
-        "loss": eval_loss,
-        "precision": precision_score(out_label_list, preds_list),
-        "recall": recall_score(out_label_list, preds_list),
-        "f1": f1_score(out_label_list, preds_list),
-    }
-
-    with open(os.path.join(args.output_dir, "test_gt.txt"), "w") as fout:
-        for lbl in out_label_list:
-            for l in lbl:
-                fout.write(l + "\t")
-            fout.write("\n")
-    with open(os.path.join(args.output_dir, "test_pred.txt"), "w") as fout:
-        for lbl in preds_list:
-            for l in lbl:
-                fout.write(l + "\t")
-            fout.write("\n")
-
-    report = classification_report(out_label_list, preds_list)
-    logger.info("\n" + report)
-
-    logger.info("***** Eval results %s *****", prefix)
-    for key in sorted(results.keys()):
-        logger.info("  %s = %s", key, str(results[key]))
-
-    return results, preds_list
-
-
 if __name__ == "__main__":
     args = parse_args()
     train(args)

ppstructure/vqa/utils.py

@@ -25,8 +25,14 @@ import paddle
 from PIL import Image, ImageDraw, ImageFont
 
 
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    paddle.seed(seed)
+
+
 def get_bio_label_maps(label_map_path):
-    with open(label_map_path, "r") as fin:
+    with open(label_map_path, "r", encoding='utf-8') as fin:
         lines = fin.readlines()
     lines = [line.strip() for line in lines]
     if "O" not in lines:
@@ -344,6 +350,8 @@ def parse_args():
     # yapf: disable
     parser.add_argument("--model_name_or_path",
                         default=None, type=str, required=True,)
+    parser.add_argument("--ser_model_type",
+                        default='LayoutXLM', type=str)
     parser.add_argument("--re_model_name_or_path",
                         default=None, type=str, required=False,)
     parser.add_argument("--train_data_dir", default=None,
@@ -357,6 +365,7 @@ def parse_args():
     parser.add_argument("--output_dir", default=None, type=str, required=True,)
     parser.add_argument("--max_seq_length", default=512, type=int,)
     parser.add_argument("--evaluate_during_training", action="store_true",)
+    parser.add_argument("--num_workers", default=8, type=int,)
     parser.add_argument("--per_gpu_train_batch_size", default=8,
                         type=int, help="Batch size per GPU/CPU for training.",)
     parser.add_argument("--per_gpu_eval_batch_size", default=8,
@@ -375,16 +384,15 @@ def parse_args():
                         help="Linear warmup over warmup_steps.",)
     parser.add_argument("--eval_steps", type=int, default=10,
                         help="eval every X updates steps.",)
-    parser.add_argument("--save_steps", type=int, default=50,
-                        help="Save checkpoint every X updates steps.",)
     parser.add_argument("--seed", type=int, default=2048,
                         help="random seed for initialization",)
 
-    parser.add_argument("--ocr_rec_model_dir", default=None, type=str, )
-    parser.add_argument("--ocr_det_model_dir", default=None, type=str, )
+    parser.add_argument("--rec_model_dir", default=None, type=str, )
+    parser.add_argument("--det_model_dir", default=None, type=str, )
     parser.add_argument(
         "--label_map_path", default="./labels/labels_ser.txt", type=str, required=False, )
     parser.add_argument("--infer_imgs", default=None, type=str, required=False)
+    parser.add_argument("--resume", action='store_true')
     parser.add_argument("--ocr_json_path", default=None,
                         type=str, required=False, help="ocr prediction results")
     # yapf: enable

ppstructure/vqa/xfun.py

@@ -79,14 +79,36 @@ class XFUNDataset(Dataset):
 
         self.entities_labels = {'HEADER': 0, 'QUESTION': 1, 'ANSWER': 2}
         self.return_keys = {
-            'bbox': 'np',
-            'input_ids': 'np',
-            'labels': 'np',
-            'attention_mask': 'np',
-            'image': 'np',
-            'token_type_ids': 'np',
-            'entities': 'dict',
-            'relations': 'dict',
+            'bbox': {
+                'type': 'np',
+                'dtype': 'int64'
+            },
+            'input_ids': {
+                'type': 'np',
+                'dtype': 'int64'
+            },
+            'labels': {
+                'type': 'np',
+                'dtype': 'int64'
+            },
+            'attention_mask': {
+                'type': 'np',
+                'dtype': 'int64'
+            },
+            'image': {
+                'type': 'np',
+                'dtype': 'float32'
+            },
+            'token_type_ids': {
+                'type': 'np',
+                'dtype': 'int64'
+            },
+            'entities': {
+                'type': 'dict'
+            },
+            'relations': {
+                'type': 'dict'
+            }
         }
 
         if load_mode == "all":
@@ -162,7 +184,7 @@ class XFUNDataset(Dataset):
         return encoded_inputs
 
     def read_all_lines(self, ):
-        with open(self.label_path, "r") as fin:
+        with open(self.label_path, "r", encoding='utf-8') as fin:
             lines = fin.readlines()
         return lines
 
@@ -412,8 +434,8 @@ class XFUNDataset(Dataset):
         return_data = {}
         for k, v in data.items():
             if k in self.return_keys:
-                if self.return_keys[k] == 'np':
-                    v = np.array(v)
+                if self.return_keys[k]['type'] == 'np':
+                    v = np.array(v, dtype=self.return_keys[k]['dtype'])
                 return_data[k] = v
         return return_data
 

requirements.txt

@@ -13,3 +13,4 @@ lxml
 premailer
 openpyxl
 fasttext==0.9.1
+

test_tipc/configs/ch_PP-OCRv2_det/train_infer_python.txt

@@ -12,9 +12,9 @@ train_model_name:latest
 train_infer_img_dir:./train_data/icdar2015/text_localization/ch4_test_images/
 null:null
 ##
-trainer:norm_train|pact_train
+trainer:norm_train
 norm_train:tools/train.py -c configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_cml.yml -o 
-pact_train:deploy/slim/quantization/quant.py -c configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_cml.yml -o
+pact_train:null
 fpgm_train:null
 distill_train:null
 null:null
@@ -26,9 +26,9 @@ null:null
 ##
 ===========================infer_params===========================
 Global.save_inference_dir:./output/
-Global.pretrained_model:
+Global.checkpoints:
 norm_export:tools/export_model.py -c configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_cml.yml -o 
-quant_export:deploy/slim/quantization/export_model.py -c configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_cml.yml -o 
+quant_export:null
 fpgm_export: 
 distill_export:null
 export1:null

test_tipc/configs/ch_PP-OCRv2_det_PACT/train_infer_python.txt

 ===========================train_params===========================
-model_name:PPOCRv2_det
+model_name:ch_PPOCRv2_det_PACT
 python:python3.7
 gpu_list:0|0,1
 Global.use_gpu:True|True