add vqa code (#7096)

* add vqa code

* add order ocr info

* rename tb-yx order

* polish configs

* add trt offline-tuning

* fix seed and remove unused configs

ppocr/data/imaug/label_ops.py

@@ -24,6 +24,7 @@ from shapely.geometry import LineString, Point, Polygon
 import json
 import copy
 from ppocr.utils.logging import get_logger
+from ppocr.data.imaug.vqa.augment import order_by_tbyx
 
 
 class ClsLabelEncode(object):
@@ -870,6 +871,7 @@ class VQATokenLabelEncode(object):
                  add_special_ids=False,
                  algorithm='LayoutXLM',
                  use_textline_bbox_info=True,
+                 order_method=None,
                  infer_mode=False,
                  ocr_engine=None,
                  **kwargs):
@@ -899,6 +901,8 @@ class VQATokenLabelEncode(object):
         self.infer_mode = infer_mode
         self.ocr_engine = ocr_engine
         self.use_textline_bbox_info = use_textline_bbox_info
+        self.order_method = order_method
+        assert self.order_method in [None, "tb-yx"]
 
     def split_bbox(self, bbox, text, tokenizer):
         words = text.split()
@@ -938,6 +942,14 @@ class VQATokenLabelEncode(object):
         # load bbox and label info
         ocr_info = self._load_ocr_info(data)
 
+        for idx in range(len(ocr_info)):
+            if "bbox" not in ocr_info[idx]:
+                ocr_info[idx]["bbox"] = self.trans_poly_to_bbox(ocr_info[idx][
+                    "points"])
+
+        if self.order_method == "tb-yx":
+            ocr_info = order_by_tbyx(ocr_info)
+
         # for re
         train_re = self.contains_re and not self.infer_mode
         if train_re:
@@ -1052,10 +1064,10 @@ class VQATokenLabelEncode(object):
         return data
 
     def trans_poly_to_bbox(self, poly):
-        x1 = np.min([p[0] for p in poly])
-        x2 = np.max([p[0] for p in poly])
-        y1 = np.min([p[1] for p in poly])
-        y2 = np.max([p[1] for p in poly])
+        x1 = int(np.min([p[0] for p in poly]))
+        x2 = int(np.max([p[0] for p in poly]))
+        y1 = int(np.min([p[1] for p in poly]))
+        y2 = int(np.max([p[1] for p in poly]))
         return [x1, y1, x2, y2]
 
     def _load_ocr_info(self, data):

ppocr/data/imaug/vqa/__init__.py

@@ -13,12 +13,10 @@
 # limitations under the License.
 
 from .token import VQATokenPad, VQASerTokenChunk, VQAReTokenChunk, VQAReTokenRelation
-from .augment import DistortBBox
 
 __all__ = [
     'VQATokenPad',
     'VQASerTokenChunk',
     'VQAReTokenChunk',
     'VQAReTokenRelation',
-    'DistortBBox',
 ]

ppocr/data/imaug/vqa/augment.py

@@ -16,22 +16,18 @@ import os
 import sys
 import numpy as np
 import random
+from copy import deepcopy
 
 
-class DistortBBox:
-    def __init__(self, prob=0.5, max_scale=1, **kwargs):
-        """Random distort bbox
-        """
-        self.prob = prob
-        self.max_scale = max_scale
-
-    def __call__(self, data):
-        if random.random() > self.prob:
-            return data
-        bbox = np.array(data['bbox'])
-        rnd_scale = (np.random.rand(*bbox.shape) - 0.5) * 2 * self.max_scale
-        bbox = np.round(bbox + rnd_scale).astype(bbox.dtype)
-        data['bbox'] = np.clip(data['bbox'], 0, 1000)
-        data['bbox'] = bbox.tolist()
-        sys.stdout.flush()
-        return data
+def order_by_tbyx(ocr_info):
+    res = sorted(ocr_info, key=lambda r: (r["bbox"][1], r["bbox"][0]))
+    for i in range(len(res) - 1):
+        for j in range(i, 0, -1):
+            if abs(res[j + 1]["bbox"][1] - res[j]["bbox"][1]) < 20 and \
+                    (res[j + 1]["bbox"][0] < res[j]["bbox"][0]):
+                tmp = deepcopy(res[j])
+                res[j] = deepcopy(res[j + 1])
+                res[j + 1] = deepcopy(tmp)
+            else:
+                break
+    return res

ppocr/losses/basic_loss.py

@@ -63,18 +63,21 @@ class KLJSLoss(object):
     def __call__(self, p1, p2, reduction="mean"):
 
         if self.mode.lower() == 'kl':
-            loss = paddle.multiply(p2, paddle.log((p2 + 1e-5) / (p1 + 1e-5) + 1e-5))
+            loss = paddle.multiply(p2,
+                                   paddle.log((p2 + 1e-5) / (p1 + 1e-5) + 1e-5))
             loss += paddle.multiply(
-                    p1, paddle.log((p1 + 1e-5) / (p2 + 1e-5) + 1e-5))
+                p1, paddle.log((p1 + 1e-5) / (p2 + 1e-5) + 1e-5))
             loss *= 0.5
         elif self.mode.lower() == "js":
-            loss = paddle.multiply(p2, paddle.log((2*p2 + 1e-5) / (p1 + p2 + 1e-5) + 1e-5))
+            loss = paddle.multiply(
+                p2, paddle.log((2 * p2 + 1e-5) / (p1 + p2 + 1e-5) + 1e-5))
             loss += paddle.multiply(
-                    p1, paddle.log((2*p1 + 1e-5) / (p1 + p2 + 1e-5) + 1e-5))
+                p1, paddle.log((2 * p1 + 1e-5) / (p1 + p2 + 1e-5) + 1e-5))
             loss *= 0.5
         else:
-            raise ValueError("The mode.lower() if KLJSLoss should be one of ['kl', 'js']")
-            
+            raise ValueError(
+                "The mode.lower() if KLJSLoss should be one of ['kl', 'js']")
+
         if reduction == "mean":
             loss = paddle.mean(loss, axis=[1, 2])
         elif reduction == "none" or reduction is None:
@@ -154,7 +157,9 @@ class LossFromOutput(nn.Layer):
         self.reduction = reduction
 
     def forward(self, predicts, batch):
-        loss = predicts[self.key]
+        loss = predicts
+        if self.key is not None and isinstance(predicts, dict):
+            loss = loss[self.key]
         if self.reduction == 'mean':
             loss = paddle.mean(loss)
         elif self.reduction == 'sum':

ppocr/losses/combined_loss.py

@@ -24,6 +24,9 @@ from .distillation_loss import DistillationCTCLoss
 from .distillation_loss import DistillationSARLoss
 from .distillation_loss import DistillationDMLLoss
 from .distillation_loss import DistillationDistanceLoss, DistillationDBLoss, DistillationDilaDBLoss
+from .distillation_loss import DistillationVQASerTokenLayoutLMLoss, DistillationSERDMLLoss
+from .distillation_loss import DistillationLossFromOutput
+from .distillation_loss import DistillationVQADistanceLoss
 
 
 class CombinedLoss(nn.Layer):

ppocr/losses/distillation_loss.py

@@ -21,8 +21,10 @@ from .rec_ctc_loss import CTCLoss
 from .rec_sar_loss import SARLoss
 from .basic_loss import DMLLoss
 from .basic_loss import DistanceLoss
+from .basic_loss import LossFromOutput
 from .det_db_loss import DBLoss
 from .det_basic_loss import BalanceLoss, MaskL1Loss, DiceLoss
+from .vqa_token_layoutlm_loss import VQASerTokenLayoutLMLoss
 
 
 def _sum_loss(loss_dict):
@@ -322,3 +324,133 @@ class DistillationDistanceLoss(DistanceLoss):
                 loss_dict["{}_{}_{}_{}".format(self.name, pair[0], pair[1],
                                                idx)] = loss
         return loss_dict
+
+
+class DistillationVQASerTokenLayoutLMLoss(VQASerTokenLayoutLMLoss):
+    def __init__(self,
+                 num_classes,
+                 model_name_list=[],
+                 key=None,
+                 name="loss_ser"):
+        super().__init__(num_classes=num_classes)
+        self.model_name_list = model_name_list
+        self.key = key
+        self.name = name
+
+    def forward(self, predicts, batch):
+        loss_dict = dict()
+        for idx, model_name in enumerate(self.model_name_list):
+            out = predicts[model_name]
+            if self.key is not None:
+                out = out[self.key]
+            loss = super().forward(out, batch)
+            loss_dict["{}_{}".format(self.name, model_name)] = loss["loss"]
+        return loss_dict
+
+
+class DistillationLossFromOutput(LossFromOutput):
+    def __init__(self,
+                 reduction="none",
+                 model_name_list=[],
+                 dist_key=None,
+                 key="loss",
+                 name="loss_re"):
+        super().__init__(key=key, reduction=reduction)
+        self.model_name_list = model_name_list
+        self.name = name
+        self.dist_key = dist_key
+
+    def forward(self, predicts, batch):
+        loss_dict = dict()
+        for idx, model_name in enumerate(self.model_name_list):
+            out = predicts[model_name]
+            if self.dist_key is not None:
+                out = out[self.dist_key]
+            loss = super().forward(out, batch)
+            loss_dict["{}_{}".format(self.name, model_name)] = loss["loss"]
+        return loss_dict
+
+
+class DistillationSERDMLLoss(DMLLoss):
+    """
+    """
+
+    def __init__(self,
+                 act="softmax",
+                 use_log=True,
+                 num_classes=7,
+                 model_name_pairs=[],
+                 key=None,
+                 name="loss_dml_ser"):
+        super().__init__(act=act, use_log=use_log)
+        assert isinstance(model_name_pairs, list)
+        self.key = key
+        self.name = name
+        self.num_classes = num_classes
+        self.model_name_pairs = model_name_pairs
+
+    def forward(self, predicts, batch):
+        loss_dict = dict()
+        for idx, pair in enumerate(self.model_name_pairs):
+            out1 = predicts[pair[0]]
+            out2 = predicts[pair[1]]
+            if self.key is not None:
+                out1 = out1[self.key]
+                out2 = out2[self.key]
+            out1 = out1.reshape([-1, out1.shape[-1]])
+            out2 = out2.reshape([-1, out2.shape[-1]])
+
+            attention_mask = batch[2]
+            if attention_mask is not None:
+                active_output = attention_mask.reshape([-1, ]) == 1
+                out1 = out1[active_output]
+                out2 = out2[active_output]
+
+            loss_dict["{}_{}".format(self.name, idx)] = super().forward(out1,
+                                                                        out2)
+
+        return loss_dict
+
+
+class DistillationVQADistanceLoss(DistanceLoss):
+    def __init__(self,
+                 mode="l2",
+                 model_name_pairs=[],
+                 key=None,
+                 name="loss_distance",
+                 **kargs):
+        super().__init__(mode=mode, **kargs)
+        assert isinstance(model_name_pairs, list)
+        self.key = key
+        self.model_name_pairs = model_name_pairs
+        self.name = name + "_l2"
+
+    def forward(self, predicts, batch):
+        loss_dict = dict()
+        for idx, pair in enumerate(self.model_name_pairs):
+            out1 = predicts[pair[0]]
+            out2 = predicts[pair[1]]
+            attention_mask = batch[2]
+            if self.key is not None:
+                out1 = out1[self.key]
+                out2 = out2[self.key]
+                if attention_mask is not None:
+                    max_len = attention_mask.shape[-1]
+                    out1 = out1[:, :max_len]
+                    out2 = out2[:, :max_len]
+                out1 = out1.reshape([-1, out1.shape[-1]])
+                out2 = out2.reshape([-1, out2.shape[-1]])
+            if attention_mask is not None:
+                active_output = attention_mask.reshape([-1, ]) == 1
+                out1 = out1[active_output]
+                out2 = out2[active_output]
+
+            loss = super().forward(out1, out2)
+            if isinstance(loss, dict):
+                for key in loss:
+                    loss_dict["{}_{}nohu_{}".format(self.name, key,
+                                                    idx)] = loss[key]
+            else:
+                loss_dict["{}_{}_{}_{}".format(self.name, pair[0], pair[1],
+                                               idx)] = loss
+        return loss_dict

ppocr/losses/vqa_token_layoutlm_loss.py

@@ -17,26 +17,30 @@ from __future__ import division
 from __future__ import print_function
 
 from paddle import nn
+from ppocr.losses.basic_loss import DMLLoss
 
 
 class VQASerTokenLayoutLMLoss(nn.Layer):
-    def __init__(self, num_classes):
+    def __init__(self, num_classes, key=None):
         super().__init__()
         self.loss_class = nn.CrossEntropyLoss()
         self.num_classes = num_classes
         self.ignore_index = self.loss_class.ignore_index
+        self.key = key
 
     def forward(self, predicts, batch):
+        if isinstance(predicts, dict) and self.key is not None:
+            predicts = predicts[self.key]
         labels = batch[5]
         attention_mask = batch[2]
         if attention_mask is not None:
             active_loss = attention_mask.reshape([-1, ]) == 1
-            active_outputs = predicts.reshape(
+            active_output = predicts.reshape(
                 [-1, self.num_classes])[active_loss]
-            active_labels = labels.reshape([-1, ])[active_loss]
-            loss = self.loss_class(active_outputs, active_labels)
+            active_label = labels.reshape([-1, ])[active_loss]
+            loss = self.loss_class(active_output, active_label)
         else:
             loss = self.loss_class(
                 predicts.reshape([-1, self.num_classes]),
                 labels.reshape([-1, ]))
-        return {'loss': loss}
+        return {'loss': loss}

ppocr/metrics/distillation_metric.py

@@ -19,6 +19,8 @@ from .rec_metric import RecMetric
 from .det_metric import DetMetric
 from .e2e_metric import E2EMetric
 from .cls_metric import ClsMetric
+from .vqa_token_ser_metric import VQASerTokenMetric
+from .vqa_token_re_metric import VQAReTokenMetric
 
 
 class DistillationMetric(object):

ppocr/modeling/architectures/base_model.py

@@ -73,28 +73,40 @@ class BaseModel(nn.Layer):
         self.return_all_feats = config.get("return_all_feats", False)
 
     def forward(self, x, data=None):
+
         y = dict()
         if self.use_transform:
             x = self.transform(x)
         x = self.backbone(x)
-        y["backbone_out"] = x
-        if self.use_neck:
-            x = self.neck(x)
-        y["neck_out"] = x
-        if self.use_head:
-            x = self.head(x, targets=data)
-        # for multi head, save ctc neck out for udml
-        if isinstance(x, dict) and 'ctc_neck' in x.keys():
-            y["neck_out"] = x["ctc_neck"]
-            y["head_out"] = x
-        elif isinstance(x, dict):
+        if isinstance(x, dict):
             y.update(x)
         else:
-            y["head_out"] = x
+            y["backbone_out"] = x
+        final_name = "backbone_out"
+        if self.use_neck:
+            x = self.neck(x)
+            if isinstance(x, dict):
+                y.update(x)
+            else:
+                y["neck_out"] = x
+            final_name = "neck_out"
+        if self.use_head:
+            x = self.head(x, targets=data)
+            # for multi head, save ctc neck out for udml
+            if isinstance(x, dict) and 'ctc_neck' in x.keys():
+                y["neck_out"] = x["ctc_neck"]
+                y["head_out"] = x
+            elif isinstance(x, dict):
+                y.update(x)
+            else:
+                y["head_out"] = x
+            final_name = "head_out"
         if self.return_all_feats:
             if self.training:
                 return y
+            elif isinstance(x, dict):
+                return x
             else:
-                return {"head_out": y["head_out"]}
+                return {final_name: x}
         else:
             return x

ppocr/modeling/backbones/vqa_layoutlm.py

@@ -22,13 +22,22 @@ from paddle import nn
 from paddlenlp.transformers import LayoutXLMModel, LayoutXLMForTokenClassification, LayoutXLMForRelationExtraction
 from paddlenlp.transformers import LayoutLMModel, LayoutLMForTokenClassification
 from paddlenlp.transformers import LayoutLMv2Model, LayoutLMv2ForTokenClassification, LayoutLMv2ForRelationExtraction
+from paddlenlp.transformers import AutoModel
 
-__all__ = ["LayoutXLMForSer", 'LayoutLMForSer']
+__all__ = ["LayoutXLMForSer", "LayoutLMForSer"]
 
 pretrained_model_dict = {
-    LayoutXLMModel: 'layoutxlm-base-uncased',
-    LayoutLMModel: 'layoutlm-base-uncased',
-    LayoutLMv2Model: 'layoutlmv2-base-uncased'
+    LayoutXLMModel: {
+        "base": "layoutxlm-base-uncased",
+        "vi": "layoutxlm-wo-backbone-base-uncased",
+    },
+    LayoutLMModel: {
+        "base": "layoutlm-base-uncased",
+    },
+    LayoutLMv2Model: {
+        "base": "layoutlmv2-base-uncased",
+        "vi": "layoutlmv2-wo-backbone-base-uncased",
+    },
 }
 
 
@@ -36,42 +45,47 @@ class NLPBaseModel(nn.Layer):
     def __init__(self,
                  base_model_class,
                  model_class,
-                 type='ser',
+                 mode="base",
+                 type="ser",
                  pretrained=True,
                  checkpoints=None,
                  **kwargs):
         super(NLPBaseModel, self).__init__()
-        if checkpoints is not None:
+        if checkpoints is not None:  # load the trained model
             self.model = model_class.from_pretrained(checkpoints)
-        elif isinstance(pretrained, (str, )) and os.path.exists(pretrained):
-            self.model = model_class.from_pretrained(pretrained)
-        else:
-            pretrained_model_name = pretrained_model_dict[base_model_class]
+        else:  # load the pretrained-model
+            pretrained_model_name = pretrained_model_dict[base_model_class][
+                mode]
             if pretrained is True:
                 base_model = base_model_class.from_pretrained(
                     pretrained_model_name)
             else:
-                base_model = base_model_class(
-                    **base_model_class.pretrained_init_configuration[
-                        pretrained_model_name])
-            if type == 'ser':
+                base_model = base_model_class.from_pretrained(pretrained)
+            if type == "ser":
                 self.model = model_class(
-                    base_model, num_classes=kwargs['num_classes'], dropout=None)
+                    base_model, num_classes=kwargs["num_classes"], dropout=None)
             else:
                 self.model = model_class(base_model, dropout=None)
         self.out_channels = 1
+        self.use_visual_backbone = True
 
 
 class LayoutLMForSer(NLPBaseModel):
-    def __init__(self, num_classes, pretrained=True, checkpoints=None,
+    def __init__(self,
+                 num_classes,
+                 pretrained=True,
+                 checkpoints=None,
+                 mode="base",
                  **kwargs):
         super(LayoutLMForSer, self).__init__(
             LayoutLMModel,
             LayoutLMForTokenClassification,
-            'ser',
+            mode,
+            "ser",
             pretrained,
             checkpoints,
-            num_classes=num_classes)
+            num_classes=num_classes, )
+        self.use_visual_backbone = False
 
     def forward(self, x):
         x = self.model(
@@ -85,62 +99,92 @@ class LayoutLMForSer(NLPBaseModel):
 
 
 class LayoutLMv2ForSer(NLPBaseModel):
-    def __init__(self, num_classes, pretrained=True, checkpoints=None,
+    def __init__(self,
+                 num_classes,
+                 pretrained=True,
+                 checkpoints=None,
+                 mode="base",
                  **kwargs):
         super(LayoutLMv2ForSer, self).__init__(
             LayoutLMv2Model,
             LayoutLMv2ForTokenClassification,
-            'ser',
+            mode,
+            "ser",
             pretrained,
             checkpoints,
             num_classes=num_classes)
+        self.use_visual_backbone = True
+        if hasattr(self.model.layoutlmv2, "use_visual_backbone"
+                   ) and self.model.layoutlmv2.use_visual_backbone is False:
+            self.use_visual_backbone = False
 
     def forward(self, x):
+        if self.use_visual_backbone is True:
+            image = x[4]
+        else:
+            image = None
         x = self.model(
             input_ids=x[0],
             bbox=x[1],
             attention_mask=x[2],
             token_type_ids=x[3],
-            image=x[4],
+            image=image,
             position_ids=None,
             head_mask=None,
             labels=None)
-        if not self.training:
+        if self.training:
+            res = {"backbone_out": x[0]}
+            res.update(x[1])
+            return res
+        else:
             return x
-        return x[0]
 
 
 class LayoutXLMForSer(NLPBaseModel):
-    def __init__(self, num_classes, pretrained=True, checkpoints=None,
+    def __init__(self,
+                 num_classes,
+                 pretrained=True,
+                 checkpoints=None,
+                 mode="base",
                  **kwargs):
         super(LayoutXLMForSer, self).__init__(
             LayoutXLMModel,
             LayoutXLMForTokenClassification,
-            'ser',
+            mode,
+            "ser",
             pretrained,
             checkpoints,
             num_classes=num_classes)
+        self.use_visual_backbone = True
 
     def forward(self, x):
+        if self.use_visual_backbone is True:
+            image = x[4]
+        else:
+            image = None
         x = self.model(
             input_ids=x[0],
             bbox=x[1],
             attention_mask=x[2],
             token_type_ids=x[3],
-            image=x[4],
+            image=image,
             position_ids=None,
             head_mask=None,
             labels=None)
-        if not self.training:
+        if self.training:
+            res = {"backbone_out": x[0]}
+            res.update(x[1])
+            return res
+        else:
             return x
-        return x[0]
 
 
 class LayoutLMv2ForRe(NLPBaseModel):
-    def __init__(self, pretrained=True, checkpoints=None, **kwargs):
-        super(LayoutLMv2ForRe, self).__init__(LayoutLMv2Model,
-                                              LayoutLMv2ForRelationExtraction,
-                                              're', pretrained, checkpoints)
+    def __init__(self, pretrained=True, checkpoints=None, mode="base",
+                 **kwargs):
+        super(LayoutLMv2ForRe, self).__init__(
+            LayoutLMv2Model, LayoutLMv2ForRelationExtraction, mode, "re",
+            pretrained, checkpoints)
 
     def forward(self, x):
         x = self.model(
@@ -158,18 +202,27 @@ class LayoutLMv2ForRe(NLPBaseModel):
 
 
 class LayoutXLMForRe(NLPBaseModel):
-    def __init__(self, pretrained=True, checkpoints=None, **kwargs):
-        super(LayoutXLMForRe, self).__init__(LayoutXLMModel,
-                                             LayoutXLMForRelationExtraction,
-                                             're', pretrained, checkpoints)
+    def __init__(self, pretrained=True, checkpoints=None, mode="base",
+                 **kwargs):
+        super(LayoutXLMForRe, self).__init__(
+            LayoutXLMModel, LayoutXLMForRelationExtraction, mode, "re",
+            pretrained, checkpoints)
+        self.use_visual_backbone = True
+        if hasattr(self.model.layoutxlm, "use_visual_backbone"
+                   ) and self.model.layoutxlm.use_visual_backbone is False:
+            self.use_visual_backbone = False
 
     def forward(self, x):
+        if self.use_visual_backbone is True:
+            image = x[4]
+        else:
+            image = None
         x = self.model(
             input_ids=x[0],
             bbox=x[1],
             attention_mask=x[2],
             token_type_ids=x[3],
-            image=x[4],
+            image=image,
             position_ids=None,
             head_mask=None,
             labels=None,

ppocr/postprocess/__init__.py

@@ -31,21 +31,38 @@ from .rec_postprocess import CTCLabelDecode, AttnLabelDecode, SRNLabelDecode, \
     SPINLabelDecode
 from .cls_postprocess import ClsPostProcess
 from .pg_postprocess import PGPostProcess
-from .vqa_token_ser_layoutlm_postprocess import VQASerTokenLayoutLMPostProcess
-from .vqa_token_re_layoutlm_postprocess import VQAReTokenLayoutLMPostProcess
+from .vqa_token_ser_layoutlm_postprocess import VQASerTokenLayoutLMPostProcess, DistillationSerPostProcess
+from .vqa_token_re_layoutlm_postprocess import VQAReTokenLayoutLMPostProcess, DistillationRePostProcess
 from .table_postprocess import TableMasterLabelDecode, TableLabelDecode
 
 
 def build_post_process(config, global_config=None):
     support_dict = [
-        'DBPostProcess', 'EASTPostProcess', 'SASTPostProcess', 'FCEPostProcess',
-        'CTCLabelDecode', 'AttnLabelDecode', 'ClsPostProcess', 'SRNLabelDecode',
-        'PGPostProcess', 'DistillationCTCLabelDecode', 'TableLabelDecode',
-        'DistillationDBPostProcess', 'NRTRLabelDecode', 'SARLabelDecode',
-        'SEEDLabelDecode', 'VQASerTokenLayoutLMPostProcess',
-        'VQAReTokenLayoutLMPostProcess', 'PRENLabelDecode',
-        'DistillationSARLabelDecode', 'ViTSTRLabelDecode', 'ABINetLabelDecode',
-        'TableMasterLabelDecode', 'SPINLabelDecode'
+        'DBPostProcess',
+        'EASTPostProcess',
+        'SASTPostProcess',
+        'FCEPostProcess',
+        'CTCLabelDecode',
+        'AttnLabelDecode',
+        'ClsPostProcess',
+        'SRNLabelDecode',
+        'PGPostProcess',
+        'DistillationCTCLabelDecode',
+        'TableLabelDecode',
+        'DistillationDBPostProcess',
+        'NRTRLabelDecode',
+        'SARLabelDecode',
+        'SEEDLabelDecode',
+        'VQASerTokenLayoutLMPostProcess',
+        'VQAReTokenLayoutLMPostProcess',
+        'PRENLabelDecode',
+        'DistillationSARLabelDecode',
+        'ViTSTRLabelDecode',
+        'ABINetLabelDecode',
+        'TableMasterLabelDecode',
+        'SPINLabelDecode',
+        'DistillationSerPostProcess',
+        'DistillationRePostProcess',
     ]
 
     if config['name'] == 'PSEPostProcess':

ppocr/postprocess/vqa_token_re_layoutlm_postprocess.py

@@ -49,3 +49,25 @@ class VQAReTokenLayoutLMPostProcess(object):
                 result.append((ocr_info_head, ocr_info_tail))
             results.append(result)
         return results
+
+
+class DistillationRePostProcess(VQAReTokenLayoutLMPostProcess):
+    """
+    DistillationRePostProcess
+    """
+
+    def __init__(self, model_name=["Student"], key=None, **kwargs):
+        super().__init__(**kwargs)
+        if not isinstance(model_name, list):
+            model_name = [model_name]
+        self.model_name = model_name
+        self.key = key
+
+    def __call__(self, preds, *args, **kwargs):
+        output = dict()
+        for name in self.model_name:
+            pred = preds[name]
+            if self.key is not None:
+                pred = pred[self.key]
+            output[name] = super().__call__(pred, *args, **kwargs)
+        return output

ppocr/postprocess/vqa_token_ser_layoutlm_postprocess.py

@@ -93,3 +93,25 @@ class VQASerTokenLayoutLMPostProcess(object):
                 ocr_info[idx]["pred"] = self.id2label_map_for_show[int(pred_id)]
             results.append(ocr_info)
         return results
+
+
+class DistillationSerPostProcess(VQASerTokenLayoutLMPostProcess):
+    """
+    DistillationSerPostProcess
+    """
+
+    def __init__(self, class_path, model_name=["Student"], key=None, **kwargs):
+        super().__init__(class_path, **kwargs)
+        if not isinstance(model_name, list):
+            model_name = [model_name]
+        self.model_name = model_name
+        self.key = key
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+        output = dict()
+        for name in self.model_name:
+            pred = preds[name]
+            if self.key is not None:
+                pred = pred[self.key]
+            output[name] = super().__call__(pred, batch=batch, *args, **kwargs)
+        return output

ppocr/utils/save_load.py

@@ -55,6 +55,9 @@ def load_model(config, model, optimizer=None, model_type='det'):
     best_model_dict = {}
 
     if model_type == 'vqa':
+        # NOTE: for vqa model, resume training is not supported now
+        if config["Architecture"]["algorithm"] in ["Distillation"]:
+            return best_model_dict
         checkpoints = config['Architecture']['Backbone']['checkpoints']
         # load vqa method metric
         if checkpoints:
@@ -78,6 +81,7 @@ def load_model(config, model, optimizer=None, model_type='det'):
                     logger.warning(
                         "{}.pdopt is not exists, params of optimizer is not loaded".
                         format(checkpoints))
+
         return best_model_dict
 
     if checkpoints:
@@ -166,15 +170,19 @@ def save_model(model,
     """
     _mkdir_if_not_exist(model_path, logger)
     model_prefix = os.path.join(model_path, prefix)
-    paddle.save(optimizer.state_dict(), model_prefix + '.pdopt')
+    if config['Architecture']["model_type"] != 'vqa':
+        paddle.save(optimizer.state_dict(), model_prefix + '.pdopt')
     if config['Architecture']["model_type"] != 'vqa':
         paddle.save(model.state_dict(), model_prefix + '.pdparams')
         metric_prefix = model_prefix
-    else:
+    else:  # for vqa system, we follow the save/load rules in NLP
         if config['Global']['distributed']:
-            model._layers.backbone.model.save_pretrained(model_prefix)
+            arch = model._layers
         else:
-            model.backbone.model.save_pretrained(model_prefix)
+            arch = model
+        if config["Architecture"]["algorithm"] in ["Distillation"]:
+            arch = arch.Student
+        arch.backbone.model.save_pretrained(model_prefix)
         metric_prefix = os.path.join(model_prefix, 'metric')
     # save metric and config
     with open(metric_prefix + '.states', 'wb') as f: