openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/losses/__init__.py

@@ -0,0 +1,62 @@
+import copy
+
+from torch import nn
+
+from .abinet_loss import ABINetLoss
+from .ar_loss import ARLoss
+from .cdistnet_loss import CDistNetLoss
+from .ce_loss import CELoss
+from .cppd_loss import CPPDLoss
+from .ctc_loss import CTCLoss
+from .igtr_loss import IGTRLoss
+from .lister_loss import LISTERLoss
+from .lpv_loss import LPVLoss
+from .mgp_loss import MGPLoss
+from .parseq_loss import PARSeqLoss
+from .robustscanner_loss import RobustScannerLoss
+from .smtr_loss import SMTRLoss
+from .srn_loss import SRNLoss
+from .visionlan_loss import VisionLANLoss
+from .cam_loss import CAMLoss
+from .seed_loss import SEEDLoss
+
+support_dict = [
+    'CTCLoss', 'ARLoss', 'CELoss', 'CPPDLoss', 'ABINetLoss', 'CDistNetLoss',
+    'VisionLANLoss', 'PARSeqLoss', 'IGTRLoss', 'SMTRLoss', 'LPVLoss',
+    'RobustScannerLoss', 'SRNLoss', 'LISTERLoss', 'GTCLoss', 'MGPLoss',
+    'CAMLoss', 'SEEDLoss'
+]
+
+
+def build_loss(config):
+    config = copy.deepcopy(config)
+    module_name = config.pop('name')
+    assert module_name in support_dict, Exception(
+        'loss only support {}'.format(support_dict))
+    module_class = eval(module_name)(**config)
+    return module_class
+
+
+class GTCLoss(nn.Module):
+
+    def __init__(self,
+                 gtc_loss,
+                 gtc_weight=1.0,
+                 ctc_weight=1.0,
+                 zero_infinity=True,
+                 **kwargs):
+        super(GTCLoss, self).__init__()
+        self.ctc_loss = CTCLoss(zero_infinity=zero_infinity)
+        self.gtc_loss = build_loss(gtc_loss)
+        self.gtc_weight = gtc_weight
+        self.ctc_weight = ctc_weight
+
+    def forward(self, predicts, batch):
+        ctc_loss = self.ctc_loss(predicts['ctc_pred'],
+                                 [None] + batch[-2:])['loss']
+        gtc_loss = self.gtc_loss(predicts['gtc_pred'], batch[:-2])['loss']
+        return {
+            'loss': self.ctc_weight * ctc_loss + self.gtc_weight * gtc_loss,
+            'ctc_loss': ctc_loss,
+            'gtc_loss': gtc_loss
+        }

openocr/openrec/losses/abinet_loss.py

@@ -0,0 +1,42 @@
+import torch
+from torch import nn
+
+
+class ABINetLoss(nn.Module):
+
+    def __init__(self,
+                 smoothing=False,
+                 ignore_index=100,
+                 align_weight=1.0,
+                 **kwargs):
+        super(ABINetLoss, self).__init__()
+        if ignore_index >= 0:
+            self.loss_func = nn.CrossEntropyLoss(reduction='mean',
+                                                 ignore_index=ignore_index)
+        else:
+            self.loss_func = nn.CrossEntropyLoss(reduction='mean')
+        self.smoothing = smoothing
+        self.align_weight = align_weight
+
+    def forward(self, pred, batch):
+        loss = {}
+        loss_sum = []
+        for name, logits in pred.items():
+            if isinstance(logits, list):
+                logit_num = len(logits)
+                if logit_num > 0:
+                    all_tgt = torch.cat([batch[1]] * logit_num, 0)
+                    all_logits = torch.cat(logits, 0)
+                    flt_logtis = all_logits.reshape([-1, all_logits.shape[2]])
+                    flt_tgt = all_tgt.reshape([-1])
+                else:
+                    continue
+            else:
+                flt_logtis = logits.reshape([-1, logits.shape[2]])
+                flt_tgt = batch[1].reshape([-1])
+
+            loss[name + '_loss'] = self.loss_func(flt_logtis, flt_tgt) * (
+                self.align_weight if name == 'align' else 1.0)
+            loss_sum.append(loss[name + '_loss'])
+        loss['loss'] = sum(loss_sum)
+        return loss

openocr/openrec/losses/ar_loss.py

@@ -0,0 +1,23 @@
+import torch.nn.functional as F
+from torch import nn
+
+
+class ARLoss(nn.Module):
+
+    def __init__(self, label_smoothing=0.1, ignore_index=0, **kwargs):
+        super(ARLoss, self).__init__()
+        self.label_smoothing = label_smoothing
+
+    def forward(self, pred, batch):
+        max_len = batch[2].max()
+        tgt = batch[1][:, 1:2 + max_len]
+        pred = pred.flatten(0, 1)
+        tgt = tgt.reshape([-1])
+        loss = F.cross_entropy(
+            pred,
+            tgt,
+            reduction='mean',
+            label_smoothing=self.label_smoothing,
+            ignore_index=pred.shape[1] + 1,
+        )  # self.loss_func(pred, tgt)
+        return {'loss': loss}

openocr/openrec/losses/cam_loss.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn.functional as F
+
+from .ar_loss import ARLoss
+
+
+def BanlanceMultiClassCrossEntropyLoss(x_o, x_t):
+    # [B, num_cls, H, W]
+    B, num_cls, H, W = x_o.shape
+    x_o = x_o.reshape(B, num_cls, H * W).permute(0, 2, 1)
+    # [B, H, W, num_cls]
+    # generate gt
+    x_t[x_t > 0.5] = 1
+    x_t[x_t <= 0.5] = 0
+    fg_x_t = x_t.sum(-1)  # [B, H, W]
+    x_t = x_t.argmax(-1)  # [B, H, W]
+    x_t[fg_x_t == 0] = num_cls - 1  # background
+    x_t = x_t.reshape(B, H * W)
+    # loss
+    weight = torch.ones((B, num_cls)).type_as(x_o)  # the weight of bg is 1.
+    num_bg = (x_t == (num_cls - 1)).sum(-1)  # [B]
+    weight[:, :-1] = (num_bg / (H * W - num_bg + 1e-5)).unsqueeze(-1).expand(
+        -1, num_cls - 1)
+    logit = F.log_softmax(x_o, dim=-1)  # [B, H*W, num_cls]
+    logit = logit * weight.unsqueeze(1)
+    loss = -logit.gather(2, x_t.unsqueeze(-1).long())
+    return loss.mean()
+
+
+class CAMLoss(ARLoss):
+
+    def __init__(self, label_smoothing=0.1, loss_weight_binary=1.5, **kwargs):
+        super(CAMLoss, self).__init__(label_smoothing=label_smoothing)
+        self.label_smoothing = label_smoothing
+        self.loss_weight_binary = loss_weight_binary
+
+    def forward(self, pred, batch):
+        binary_mask = batch[-1]
+        rec_loss = super().forward(pred['rec_output'], batch[:-1])['loss']
+        output = pred
+        loss_binary = self.loss_weight_binary * BanlanceMultiClassCrossEntropyLoss(
+            output['pred_binary'], binary_mask)
+
+        return {
+            'loss': rec_loss + loss_binary,
+            'rec_loss': rec_loss,
+            'loss_binary': loss_binary
+        }

openocr/openrec/losses/cdistnet_loss.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class CDistNetLoss(nn.Module):
+
+    def __init__(self, smoothing=True, ignore_index=0, **kwargs):
+        super(CDistNetLoss, self).__init__()
+        if ignore_index >= 0 and not smoothing:
+            self.loss_func = nn.CrossEntropyLoss(reduction='mean',
+                                                 ignore_index=ignore_index)
+        self.smoothing = smoothing
+
+    def forward(self, pred, batch):
+        pred = pred['res']
+        tgt = batch[1][:, 1:]
+        pred = pred.reshape([-1, pred.shape[2]])
+        tgt = tgt.reshape([-1])
+        if self.smoothing:
+            eps = 0.1
+            n_class = pred.shape[1]
+            one_hot = F.one_hot(tgt.long(), num_classes=pred.shape[1])
+            torch.set_printoptions(profile='full')
+            one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
+            log_prb = F.log_softmax(pred, dim=1)
+            non_pad_mask = torch.not_equal(
+                tgt, torch.zeros(tgt.shape, dtype=tgt.dtype,
+                                 device=tgt.device))
+            loss = -(one_hot * log_prb).sum(dim=1)
+            loss = loss.masked_select(non_pad_mask).mean()
+        else:
+            loss = self.loss_func(pred, tgt)
+        return {'loss': loss}

openocr/openrec/losses/ce_loss.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class CELoss(nn.Module):
+
+    def __init__(self,
+                 smoothing=False,
+                 with_all=False,
+                 ignore_index=-1,
+                 **kwargs):
+        super(CELoss, self).__init__()
+        if ignore_index >= 0:
+            self.loss_func = nn.CrossEntropyLoss(reduction='mean',
+                                                 ignore_index=ignore_index)
+        else:
+            self.loss_func = nn.CrossEntropyLoss(reduction='mean')
+        self.smoothing = smoothing
+        self.with_all = with_all
+
+    def forward(self, pred, batch):
+        pred = pred['res']
+        if isinstance(pred, dict):  # for ABINet
+            loss = {}
+            loss_sum = []
+            for name, logits in pred.items():
+                if isinstance(logits, list):
+                    logit_num = len(logits)
+                    all_tgt = torch.cat([batch[1]] * logit_num, 0)
+                    all_logits = torch.cat(logits, 0)
+                    flt_logtis = all_logits.reshape([-1, all_logits.shape[2]])
+                    flt_tgt = all_tgt.reshape([-1])
+                else:
+                    flt_logtis = logits.reshape([-1, logits.shape[2]])
+                    flt_tgt = batch[1].reshape([-1])
+                loss[name + '_loss'] = self.loss_func(flt_logtis, flt_tgt)
+                loss_sum.append(loss[name + '_loss'])
+            loss['loss'] = sum(loss_sum)
+            return loss
+        else:
+            if self.with_all:  # for ViTSTR
+                tgt = batch[1]
+                pred = pred.reshape([-1, pred.shape[2]])
+                tgt = tgt.reshape([-1])
+                loss = self.loss_func(pred, tgt)
+                return {'loss': loss}
+            else:  # for NRTR
+                max_len = batch[2].max()
+                tgt = batch[1][:, 1:2 + max_len]
+                pred = pred.reshape([-1, pred.shape[2]])
+                tgt = tgt.reshape([-1])
+                if self.smoothing:
+                    eps = 0.1
+                    pred.shape[1]
+                    one_hot = F.one_hot(tgt, pred.shape[1])
+                    one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (-1)
+                    log_prb = F.log_softmax(pred, dim=1)
+                    non_pad_mask = torch.not_equal(
+                        tgt,
+                        torch.zeros(tgt.shape,
+                                    dtype=tgt.dtype,
+                                    device=tgt.device))
+                    loss = -(one_hot * log_prb).sum(dim=1)
+                    loss = loss.masked_select(non_pad_mask).mean()
+                else:
+                    loss = self.loss_func(pred, tgt)
+                return {'loss': loss}

openocr/openrec/losses/cppd_loss.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class CPPDLoss(nn.Module):
+
+    def __init__(self,
+                 smoothing=False,
+                 ignore_index=100,
+                 pos_len=False,
+                 sideloss_weight=1.0,
+                 max_len=25,
+                 **kwargs):
+        super(CPPDLoss, self).__init__()
+        self.edge_ce = nn.CrossEntropyLoss(reduction='mean',
+                                           ignore_index=ignore_index)
+        self.char_node_ce = nn.CrossEntropyLoss(reduction='mean')
+        if pos_len:
+            self.pos_node_ce = nn.CrossEntropyLoss(reduction='mean',
+                                                   ignore_index=ignore_index)
+        else:
+            self.pos_node_ce = nn.BCEWithLogitsLoss(reduction='mean')
+
+        self.smoothing = smoothing
+        self.ignore_index = ignore_index
+        self.pos_len = pos_len
+        self.sideloss_weight = sideloss_weight
+        self.max_len = max_len + 1
+
+    def label_smoothing_ce(self, preds, targets):
+        zeros_ = torch.zeros_like(targets)
+        ignore_index_ = zeros_ + self.ignore_index
+        non_pad_mask = torch.not_equal(targets, ignore_index_)
+
+        tgts = torch.where(targets == ignore_index_, zeros_, targets)
+        eps = 0.1
+        n_class = preds.shape[1]
+        one_hot = F.one_hot(tgts, preds.shape[1])
+        one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
+        log_prb = F.log_softmax(preds, dim=1)
+        loss = -(one_hot * log_prb).sum(dim=1)
+        loss = loss.masked_select(non_pad_mask).mean()
+        return loss
+
+    def forward(self, pred, batch):
+        node_feats, edge_feats = pred
+        node_tgt = batch[2]
+        char_tgt = batch[1]
+
+        # updated code
+        char_num_label = torch.clip(node_tgt[:, :-self.max_len].flatten(0, 1),
+                                    0, node_feats[0].shape[-1] - 1)
+        loss_char_node = self.char_node_ce(node_feats[0].flatten(0, 1),
+                                           char_num_label)
+        if self.pos_len:
+            loss_pos_node = self.pos_node_ce(
+                node_feats[1].flatten(0, 1),
+                node_tgt[:, -self.max_len:].flatten(0, 1))
+        else:
+            loss_pos_node = self.pos_node_ce(
+                node_feats[1].flatten(0, 1),
+                node_tgt[:, -self.max_len:].flatten(0, 1).float())
+        loss_node = loss_char_node + loss_pos_node
+        # -----
+        edge_feats = edge_feats.flatten(0, 1)
+        char_tgt = char_tgt.flatten(0, 1)
+        if self.smoothing:
+            loss_edge = self.label_smoothing_ce(edge_feats, char_tgt)
+        else:
+            loss_edge = self.edge_ce(edge_feats, char_tgt)
+
+        return {
+            'loss': self.sideloss_weight * loss_node + loss_edge,
+            'loss_node': self.sideloss_weight * loss_node,
+            'loss_edge': loss_edge,
+        }

openocr/openrec/losses/ctc_loss.py

@@ -0,0 +1,33 @@
+import torch
+from torch import nn
+
+
+class CTCLoss(nn.Module):
+
+    def __init__(self, use_focal_loss=False, zero_infinity=False, **kwargs):
+        super(CTCLoss, self).__init__()
+        self.loss_func = nn.CTCLoss(blank=0,
+                                    reduction='none',
+                                    zero_infinity=zero_infinity)
+        self.use_focal_loss = use_focal_loss
+
+    def forward(self, predicts, batch):
+        # predicts = predicts['res']
+
+        batch_size = predicts.size(0)
+        label, label_length = batch[1], batch[2]
+        predicts = predicts.log_softmax(2)
+        predicts = predicts.permute(1, 0, 2)
+        preds_lengths = torch.tensor([predicts.size(0)] * batch_size,
+                                     dtype=torch.long)
+        loss = self.loss_func(predicts, label, preds_lengths, label_length)
+
+        if self.use_focal_loss:
+            # Use torch.clamp to limit the range of loss, avoiding overflow in exponential calculation
+            clamped_loss = torch.clamp(loss, min=-20, max=20)
+            weight = 1 - torch.exp(-clamped_loss)
+            weight = torch.square(weight)
+            # Use torch.where to avoid multiplying by zero weight
+            loss = torch.where(weight > 0, loss * weight, loss)
+        loss = loss.mean()
+        return {'loss': loss}

openocr/openrec/losses/igtr_loss.py

@@ -0,0 +1,12 @@
+from torch import nn
+
+
+class IGTRLoss(nn.Module):
+
+    def __init__(self, **kwargs):
+        super(IGTRLoss, self).__init__()
+
+    def forward(self, predicts, batch):
+        if isinstance(predicts, list):
+            predicts = predicts[0]
+        return predicts

openocr/openrec/losses/lister_loss.py

@@ -0,0 +1,14 @@
+from torch import nn
+
+
+class LISTERLoss(nn.Module):
+
+    def __init__(self, **kwargs):
+        super(LISTERLoss, self).__init__()
+
+    def forward(self, predicts, batch):
+        # predicts = predicts['res']
+        # loss = predicts
+        if isinstance(predicts, list):
+            predicts = predicts[0]
+        return predicts

openocr/openrec/losses/lpv_loss.py

@@ -0,0 +1,30 @@
+import torch.nn.functional as F
+from torch import nn
+
+
+class LPVLoss(nn.Module):
+
+    def __init__(self, label_smoothing=0.0, **kwargs):
+        super(LPVLoss, self).__init__()
+        self.label_smoothing = label_smoothing
+
+    def forward(self, preds, batch):
+        max_len = batch[2].max()
+        tgt = batch[1][:, 1:2 + max_len]
+
+        tgt = tgt.flatten(0, 1)
+        loss = 0
+        loss_dict = {}
+        for i, pred in enumerate(preds):
+            pred = pred.flatten(0, 1)
+            loss_i = F.cross_entropy(
+                pred,
+                tgt,
+                reduction='mean',
+                label_smoothing=self.label_smoothing,
+                ignore_index=pred.shape[1] + 1,
+            )  # self.loss_func(pred, tgt)
+            loss += loss_i
+            loss_dict['loss' + str(i)] = loss_i
+        loss_dict['loss'] = loss
+        return loss_dict

openocr/openrec/losses/mgp_loss.py

@@ -0,0 +1,34 @@
+from torch import nn
+
+
+class MGPLoss(nn.Module):
+
+    def __init__(self, only_char=False, **kwargs):
+        super(MGPLoss, self).__init__()
+        self.ce = nn.CrossEntropyLoss(reduction='mean', ignore_index=0)
+        self.only_char = only_char
+
+    def forward(self, pred, batch):
+        if self.only_char:
+            char_feats = pred
+            char_tgt = batch[1].flatten(0, 1)
+            char_loss = self.ce(char_feats.flatten(0, 1), char_tgt)
+            return {'loss': char_loss}
+        else:
+            return self.forward_all(pred, batch)
+
+    def forward_all(self, pred, batch):
+        char_feats, dpe_feats, wp_feats = pred
+        char_tgt = batch[1].flatten(0, 1)
+        dpe_tgt = batch[2].flatten(0, 1)
+        wp_tgt = batch[3].flatten(0, 1)
+        char_loss = self.ce(char_feats.flatten(0, 1), char_tgt)
+        dpe_loss = self.ce(dpe_feats.flatten(0, 1), dpe_tgt)
+        wp_loss = self.ce(wp_feats.flatten(0, 1), wp_tgt)
+        loss = char_loss + dpe_loss + wp_loss
+        return {
+            'loss': loss,
+            'char_loss': char_loss,
+            'dpe_loss': dpe_loss,
+            'wp_loss': wp_loss
+        }

openocr/openrec/losses/parseq_loss.py

@@ -0,0 +1,12 @@
+from torch import nn
+
+
+class PARSeqLoss(nn.Module):
+
+    def __init__(self, **kwargs):
+        super(PARSeqLoss, self).__init__()
+
+    def forward(self, predicts, batch):
+        # predicts = predicts['res']
+        loss, _ = predicts
+        return {'loss': loss}

openocr/openrec/losses/robustscanner_loss.py

@@ -0,0 +1,20 @@
+from torch import nn
+
+
+class RobustScannerLoss(nn.Module):
+
+    def __init__(self, **kwargs):
+        super(RobustScannerLoss, self).__init__()
+        ignore_index = kwargs.get('ignore_index', 38)
+        self.loss_func = nn.CrossEntropyLoss(reduction='mean',
+                                             ignore_index=ignore_index)
+
+    def forward(self, pred, batch):
+        pred = pred[:, :-1, :]
+
+        label = batch[1][:, 1:].reshape([-1])
+
+        inputs = pred.reshape([-1, pred.shape[2]])
+
+        loss = self.loss_func(inputs, label)
+        return {'loss': loss}

openocr/openrec/losses/seed_loss.py

@@ -0,0 +1,46 @@
+import torch.nn.functional as F
+from torch import nn
+import torch
+
+
+class CosineEmbeddingLoss(nn.Module):
+
+    def __init__(self, margin=0.0):
+        super(CosineEmbeddingLoss, self).__init__()
+        self.margin = margin
+        self.epsilon = 1e-12
+
+    def forward(self, x1, x2):
+        similarity = torch.sum(x1 * x2, axis=-1) / (
+            torch.norm(x1, dim=-1) * torch.norm(x2, dim=-1) + self.epsilon)
+        return (1 - similarity).mean()
+
+
+class SEEDLoss(nn.Module):
+
+    def __init__(self, label_smoothing=0.1, ignore_index=0, **kwargs):
+        super(SEEDLoss, self).__init__()
+        self.label_smoothing = label_smoothing
+        self.loss_sem = CosineEmbeddingLoss()
+
+    def forward(self, preds, batch):
+        embedding_vectors, pred = preds
+        max_len = batch[2].max()
+        tgt = batch[1][:, 1:2 + max_len]
+        pred = pred.flatten(0, 1)
+        tgt = tgt.reshape([-1])
+        loss = F.cross_entropy(
+            pred,
+            tgt,
+            reduction='mean',
+            label_smoothing=self.label_smoothing,
+            ignore_index=pred.shape[1] + 1,
+        )  # self.loss_func(pred, tgt)
+        sem_target = batch[3].float()
+
+        sem_loss = torch.sum(self.loss_sem(embedding_vectors, sem_target))
+        return {
+            'loss': loss + 0.1 * sem_loss,
+            'rec_loss': loss,
+            'sem_loss': sem_loss
+        }

openocr/openrec/losses/smtr_loss.py

@@ -0,0 +1,12 @@
+from torch import nn
+
+
+class SMTRLoss(nn.Module):
+
+    def __init__(self, **kwargs):
+        super(SMTRLoss, self).__init__()
+
+    def forward(self, predicts, batch):
+        if isinstance(predicts, list):
+            predicts = predicts[0]
+        return predicts

openocr/openrec/losses/srn_loss.py

@@ -0,0 +1,40 @@
+import torch.nn.functional as F
+from torch import nn
+
+
+class SRNLoss(nn.Module):
+
+    def __init__(self, label_smoothing=0.0, **kwargs):
+        super(SRNLoss, self).__init__()
+        self.label_smoothing = label_smoothing
+
+    def forward(self, preds, batch):
+        pvam_preds, gsrm_preds, vsfd_preds = preds
+
+        label = batch[1].reshape([-1])
+
+        ignore_index = pvam_preds.shape[-1] + 1
+
+        loss_pvam = F.cross_entropy(pvam_preds,
+                                    label,
+                                    reduction='mean',
+                                    label_smoothing=self.label_smoothing,
+                                    ignore_index=ignore_index)
+        loss_gsrm = F.cross_entropy(gsrm_preds,
+                                    label,
+                                    reduction='mean',
+                                    label_smoothing=self.label_smoothing,
+                                    ignore_index=ignore_index)
+        loss_vsfd = F.cross_entropy(vsfd_preds,
+                                    label,
+                                    reduction='mean',
+                                    label_smoothing=self.label_smoothing,
+                                    ignore_index=ignore_index)
+
+        loss_dict = {}
+        loss_dict['loss_pvam'] = loss_pvam
+        loss_dict['loss_gsrm'] = loss_gsrm
+        loss_dict['loss_vsfd'] = loss_vsfd
+
+        loss_dict['loss'] = loss_pvam * 3.0 + loss_gsrm * 0.15 + loss_vsfd
+        return loss_dict

openocr/openrec/losses/visionlan_loss.py

@@ -0,0 +1,58 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def flatten_label(target):
+    label_flatten = []
+    label_length = []
+    for i in range(0, target.size()[0]):
+        cur_label = target[i].tolist()
+        label_flatten += cur_label[:cur_label.index(0) + 1]
+        label_length.append(cur_label.index(0) + 1)
+    label_flatten = torch.LongTensor(label_flatten)
+    label_length = torch.IntTensor(label_length)
+    return (label_flatten, label_length)
+
+
+def _flatten(sources, lengths):
+    return torch.cat([t[:l] for t, l in zip(sources, lengths)])
+
+
+class VisionLANLoss(nn.Module):
+
+    def __init__(self,
+                 training_step='LA',
+                 ratio_res=0.5,
+                 ratio_sub=0.5,
+                 **kwargs):
+        super(VisionLANLoss, self).__init__()
+        self.loss_func = nn.CrossEntropyLoss(reduction='mean')
+        self.ratio_res = ratio_res
+        self.ratio_sub = ratio_sub
+        assert training_step in ['LF_1', 'LF_2', 'LA']
+        self.training_step = training_step
+
+    def forward(self, pred, batch):
+        text_pre, text_rem, text_mas, _ = pred
+        target = batch[1].to(dtype=torch.int64)
+        label_flatten, length = flatten_label(target)
+        text_pre = _flatten(text_pre, length)
+        if self.training_step == 'LF_1':
+            loss = self.loss_func(text_pre, label_flatten.to(text_pre.device))
+        else:
+            target_res = batch[2].to(dtype=torch.int64)
+            target_sub = batch[3].to(dtype=torch.int64)
+            label_flatten_res, length_res = flatten_label(target_res)
+            label_flatten_sub, length_sub = flatten_label(target_sub)
+            text_rem = _flatten(text_rem, length_res)
+            text_mas = _flatten(text_mas, length_sub)
+            loss_ori = self.loss_func(text_pre,
+                                      label_flatten.to(text_pre.device))
+            loss_res = self.loss_func(text_rem,
+                                      label_flatten_res.to(text_rem.device))
+            loss_mas = self.loss_func(text_mas,
+                                      label_flatten_sub.to(text_mas.device))
+            loss = loss_ori + loss_res * self.ratio_res + loss_mas * self.ratio_sub
+
+        return {'loss': loss}