pyxllib 0.3.96__py3-none-any.whl → 0.3.200__py3-none-any.whl

pyxlpr/ppocr/modeling/heads/rec_sar_head.py

@@ -1,402 +0,0 @@
-# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
-#
-# 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.
-"""
-This code is refer from: 
-https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/encoders/sar_encoder.py
-https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/decoders/sar_decoder.py
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import math
-import paddle
-from paddle import ParamAttr
-import paddle.nn as nn
-import paddle.nn.functional as F
-
-
-class SAREncoder(nn.Layer):
-    """
-    Args:
-        enc_bi_rnn (bool): If True, use bidirectional RNN in encoder.
-        enc_drop_rnn (float): Dropout probability of RNN layer in encoder.
-        enc_gru (bool): If True, use GRU, else LSTM in encoder.
-        d_model (int): Dim of channels from backbone.
-        d_enc (int): Dim of encoder RNN layer.
-        mask (bool): If True, mask padding in RNN sequence.
-    """
-
-    def __init__(self,
-                 enc_bi_rnn=False,
-                 enc_drop_rnn=0.1,
-                 enc_gru=False,
-                 d_model=512,
-                 d_enc=512,
-                 mask=True,
-                 **kwargs):
-        super().__init__()
-        assert isinstance(enc_bi_rnn, bool)
-        assert isinstance(enc_drop_rnn, (int, float))
-        assert 0 <= enc_drop_rnn < 1.0
-        assert isinstance(enc_gru, bool)
-        assert isinstance(d_model, int)
-        assert isinstance(d_enc, int)
-        assert isinstance(mask, bool)
-
-        self.enc_bi_rnn = enc_bi_rnn
-        self.enc_drop_rnn = enc_drop_rnn
-        self.mask = mask
-
-        # LSTM Encoder
-        if enc_bi_rnn:
-            direction = 'bidirectional'
-        else:
-            direction = 'forward'
-        kwargs = dict(
-            input_size=d_model,
-            hidden_size=d_enc,
-            num_layers=2,
-            time_major=False,
-            dropout=enc_drop_rnn,
-            direction=direction)
-        if enc_gru:
-            self.rnn_encoder = nn.GRU(**kwargs)
-        else:
-            self.rnn_encoder = nn.LSTM(**kwargs)
-
-        # global feature transformation
-        encoder_rnn_out_size = d_enc * (int(enc_bi_rnn) + 1)
-        self.linear = nn.Linear(encoder_rnn_out_size, encoder_rnn_out_size)
-
-    def forward(self, feat, img_metas=None):
-        if img_metas is not None:
-            assert len(img_metas[0]) == feat.shape[0]
-
-        valid_ratios = None
-        if img_metas is not None and self.mask:
-            valid_ratios = img_metas[-1]
-
-        h_feat = feat.shape[2]  # bsz c h w
-        feat_v = F.max_pool2d(
-            feat, kernel_size=(h_feat, 1), stride=1, padding=0)
-        feat_v = feat_v.squeeze(2)  # bsz * C * W
-        feat_v = paddle.transpose(feat_v, perm=[0, 2, 1])  # bsz * W * C
-        holistic_feat = self.rnn_encoder(feat_v)[0]  # bsz * T * C
-
-        if valid_ratios is not None:
-            valid_hf = []
-            T = holistic_feat.shape[1]
-            for i, valid_ratio in enumerate(valid_ratios):
-                valid_step = min(T, math.ceil(T * valid_ratio)) - 1
-                valid_hf.append(holistic_feat[i, valid_step, :])
-            valid_hf = paddle.stack(valid_hf, axis=0)
-        else:
-            valid_hf = holistic_feat[:, -1, :]  # bsz * C
-        holistic_feat = self.linear(valid_hf)  # bsz * C
-
-        return holistic_feat
-
-
-class BaseDecoder(nn.Layer):
-    def __init__(self, **kwargs):
-        super().__init__()
-
-    def forward_train(self, feat, out_enc, targets, img_metas):
-        raise NotImplementedError
-
-    def forward_test(self, feat, out_enc, img_metas):
-        raise NotImplementedError
-
-    def forward(self,
-                feat,
-                out_enc,
-                label=None,
-                img_metas=None,
-                train_mode=True):
-        self.train_mode = train_mode
-
-        if train_mode:
-            return self.forward_train(feat, out_enc, label, img_metas)
-        return self.forward_test(feat, out_enc, img_metas)
-
-
-class ParallelSARDecoder(BaseDecoder):
-    """
-    Args:
-        out_channels (int): Output class number.
-        enc_bi_rnn (bool): If True, use bidirectional RNN in encoder.
-        dec_bi_rnn (bool): If True, use bidirectional RNN in decoder.
-        dec_drop_rnn (float): Dropout of RNN layer in decoder.
-        dec_gru (bool): If True, use GRU, else LSTM in decoder.
-        d_model (int): Dim of channels from backbone.
-        d_enc (int): Dim of encoder RNN layer.
-        d_k (int): Dim of channels of attention module.
-        pred_dropout (float): Dropout probability of prediction layer.
-        max_seq_len (int): Maximum sequence length for decoding.
-        mask (bool): If True, mask padding in feature map.
-        start_idx (int): Index of start token.
-        padding_idx (int): Index of padding token.
-        pred_concat (bool): If True, concat glimpse feature from
-            attention with holistic feature and hidden state.
-    """
-
-    def __init__(
-            self,
-            out_channels,  # 90 + unknown + start + padding
-            enc_bi_rnn=False,
-            dec_bi_rnn=False,
-            dec_drop_rnn=0.0,
-            dec_gru=False,
-            d_model=512,
-            d_enc=512,
-            d_k=64,
-            pred_dropout=0.1,
-            max_text_length=30,
-            mask=True,
-            pred_concat=True,
-            **kwargs):
-        super().__init__()
-
-        self.num_classes = out_channels
-        self.enc_bi_rnn = enc_bi_rnn
-        self.d_k = d_k
-        self.start_idx = out_channels - 2
-        self.padding_idx = out_channels - 1
-        self.max_seq_len = max_text_length
-        self.mask = mask
-        self.pred_concat = pred_concat
-
-        encoder_rnn_out_size = d_enc * (int(enc_bi_rnn) + 1)
-        decoder_rnn_out_size = encoder_rnn_out_size * (int(dec_bi_rnn) + 1)
-
-        # 2D attention layer
-        self.conv1x1_1 = nn.Linear(decoder_rnn_out_size, d_k)
-        self.conv3x3_1 = nn.Conv2D(
-            d_model, d_k, kernel_size=3, stride=1, padding=1)
-        self.conv1x1_2 = nn.Linear(d_k, 1)
-
-        # Decoder RNN layer
-        if dec_bi_rnn:
-            direction = 'bidirectional'
-        else:
-            direction = 'forward'
-
-        kwargs = dict(
-            input_size=encoder_rnn_out_size,
-            hidden_size=encoder_rnn_out_size,
-            num_layers=2,
-            time_major=False,
-            dropout=dec_drop_rnn,
-            direction=direction)
-        if dec_gru:
-            self.rnn_decoder = nn.GRU(**kwargs)
-        else:
-            self.rnn_decoder = nn.LSTM(**kwargs)
-
-        # Decoder input embedding
-        self.embedding = nn.Embedding(
-            self.num_classes,
-            encoder_rnn_out_size,
-            padding_idx=self.padding_idx)
-
-        # Prediction layer
-        self.pred_dropout = nn.Dropout(pred_dropout)
-        pred_num_classes = self.num_classes - 1
-        if pred_concat:
-            fc_in_channel = decoder_rnn_out_size + d_model + d_enc
-        else:
-            fc_in_channel = d_model
-        self.prediction = nn.Linear(fc_in_channel, pred_num_classes)
-
-    def _2d_attention(self,
-                      decoder_input,
-                      feat,
-                      holistic_feat,
-                      valid_ratios=None):
-
-        y = self.rnn_decoder(decoder_input)[0]
-        # y: bsz * (seq_len + 1) * hidden_size
-
-        attn_query = self.conv1x1_1(y)  # bsz * (seq_len + 1) * attn_size
-        bsz, seq_len, attn_size = attn_query.shape
-        attn_query = paddle.unsqueeze(attn_query, axis=[3, 4])
-        # (bsz, seq_len + 1, attn_size, 1, 1)
-
-        attn_key = self.conv3x3_1(feat)
-        # bsz * attn_size * h * w
-        attn_key = attn_key.unsqueeze(1)
-        # bsz * 1 * attn_size * h * w
-
-        attn_weight = paddle.tanh(paddle.add(attn_key, attn_query))
-
-        # bsz * (seq_len + 1) * attn_size * h * w
-        attn_weight = paddle.transpose(attn_weight, perm=[0, 1, 3, 4, 2])
-        # bsz * (seq_len + 1) * h * w * attn_size
-        attn_weight = self.conv1x1_2(attn_weight)
-        # bsz * (seq_len + 1) * h * w * 1
-        bsz, T, h, w, c = attn_weight.shape
-        assert c == 1
-
-        if valid_ratios is not None:
-            # cal mask of attention weight
-            for i, valid_ratio in enumerate(valid_ratios):
-                valid_width = min(w, math.ceil(w * valid_ratio))
-                if valid_width < w:
-                    attn_weight[i, :, :, valid_width:, :] = float('-inf')
-
-        attn_weight = paddle.reshape(attn_weight, [bsz, T, -1])
-        attn_weight = F.softmax(attn_weight, axis=-1)
-
-        attn_weight = paddle.reshape(attn_weight, [bsz, T, h, w, c])
-        attn_weight = paddle.transpose(attn_weight, perm=[0, 1, 4, 2, 3])
-        # attn_weight: bsz * T * c * h * w
-        # feat: bsz * c * h * w
-        attn_feat = paddle.sum(paddle.multiply(feat.unsqueeze(1), attn_weight),
-                               (3, 4),
-                               keepdim=False)
-        # bsz * (seq_len + 1) * C
-
-        # Linear transformation
-        if self.pred_concat:
-            hf_c = holistic_feat.shape[-1]
-            holistic_feat = paddle.expand(
-                holistic_feat, shape=[bsz, seq_len, hf_c])
-            y = self.prediction(paddle.concat((y, attn_feat, holistic_feat), 2))
-        else:
-            y = self.prediction(attn_feat)
-        # bsz * (seq_len + 1) * num_classes
-        if self.train_mode:
-            y = self.pred_dropout(y)
-
-        return y
-
-    def forward_train(self, feat, out_enc, label, img_metas):
-        '''
-        img_metas: [label, valid_ratio]
-        '''
-        if img_metas is not None:
-            assert len(img_metas[0]) == feat.shape[0]
-
-        valid_ratios = None
-        if img_metas is not None and self.mask:
-            valid_ratios = img_metas[-1]
-
-        lab_embedding = self.embedding(label)
-        # bsz * seq_len * emb_dim
-        out_enc = out_enc.unsqueeze(1)
-        # bsz * 1 * emb_dim
-        in_dec = paddle.concat((out_enc, lab_embedding), axis=1)
-        # bsz * (seq_len + 1) * C
-        out_dec = self._2d_attention(
-            in_dec, feat, out_enc, valid_ratios=valid_ratios)
-        # bsz * (seq_len + 1) * num_classes
-
-        return out_dec[:, 1:, :]  # bsz * seq_len * num_classes
-
-    def forward_test(self, feat, out_enc, img_metas):
-        if img_metas is not None:
-            assert len(img_metas[0]) == feat.shape[0]
-
-        valid_ratios = None
-        if img_metas is not None and self.mask:
-            valid_ratios = img_metas[-1]
-
-        seq_len = self.max_seq_len
-        bsz = feat.shape[0]
-        start_token = paddle.full(
-            (bsz, ), fill_value=self.start_idx, dtype='int64')
-        # bsz
-        start_token = self.embedding(start_token)
-        # bsz * emb_dim
-        emb_dim = start_token.shape[1]
-        start_token = start_token.unsqueeze(1)
-        start_token = paddle.expand(start_token, shape=[bsz, seq_len, emb_dim])
-        # bsz * seq_len * emb_dim
-        out_enc = out_enc.unsqueeze(1)
-        # bsz * 1 * emb_dim
-        decoder_input = paddle.concat((out_enc, start_token), axis=1)
-        # bsz * (seq_len + 1) * emb_dim
-
-        outputs = []
-        for i in range(1, seq_len + 1):
-            decoder_output = self._2d_attention(
-                decoder_input, feat, out_enc, valid_ratios=valid_ratios)
-            char_output = decoder_output[:, i, :]  # bsz * num_classes
-            char_output = F.softmax(char_output, -1)
-            outputs.append(char_output)
-            max_idx = paddle.argmax(char_output, axis=1, keepdim=False)
-            char_embedding = self.embedding(max_idx)  # bsz * emb_dim
-            if i < seq_len:
-                decoder_input[:, i + 1, :] = char_embedding
-
-        outputs = paddle.stack(outputs, 1)  # bsz * seq_len * num_classes
-
-        return outputs
-
-
-class SARHead(nn.Layer):
-    def __init__(self,
-                 out_channels,
-                 enc_bi_rnn=False,
-                 enc_drop_rnn=0.1,
-                 enc_gru=False,
-                 dec_bi_rnn=False,
-                 dec_drop_rnn=0.0,
-                 dec_gru=False,
-                 d_k=512,
-                 pred_dropout=0.1,
-                 max_text_length=30,
-                 pred_concat=True,
-                 **kwargs):
-        super(SARHead, self).__init__()
-
-        # encoder module
-        self.encoder = SAREncoder(
-            enc_bi_rnn=enc_bi_rnn, enc_drop_rnn=enc_drop_rnn, enc_gru=enc_gru)
-
-        # decoder module
-        self.decoder = ParallelSARDecoder(
-            out_channels=out_channels,
-            enc_bi_rnn=enc_bi_rnn,
-            dec_bi_rnn=dec_bi_rnn,
-            dec_drop_rnn=dec_drop_rnn,
-            dec_gru=dec_gru,
-            d_k=d_k,
-            pred_dropout=pred_dropout,
-            max_text_length=max_text_length,
-            pred_concat=pred_concat)
-
-    def forward(self, feat, targets=None):
-        '''
-        img_metas: [label, valid_ratio]
-        '''
-        holistic_feat = self.encoder(feat, targets)  # bsz c
-
-        if self.training:
-            label = targets[0]  # label
-            label = paddle.to_tensor(label, dtype='int64')
-            final_out = self.decoder(
-                feat, holistic_feat, label, img_metas=targets)
-        if not self.training:
-            final_out = self.decoder(
-                feat,
-                holistic_feat,
-                label=None,
-                img_metas=targets,
-                train_mode=False)
-            # (bsz, seq_len, num_classes)
-
-        return final_out

pyxlpr/ppocr/modeling/heads/rec_srn_head.py

@@ -1,280 +0,0 @@
-# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
-#
-# 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 __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import math
-import paddle
-from paddle import nn, ParamAttr
-from paddle.nn import functional as F
-import paddle.fluid as fluid
-import numpy as np
-from .self_attention import WrapEncoderForFeature
-from .self_attention import WrapEncoder
-from paddle.static import Program
-from pyxlpr.ppocr.modeling.backbones.rec_resnet_fpn import ResNetFPN
-import paddle.fluid.framework as framework
-
-from collections import OrderedDict
-gradient_clip = 10
-
-
-class PVAM(nn.Layer):
-    def __init__(self, in_channels, char_num, max_text_length, num_heads,
-                 num_encoder_tus, hidden_dims):
-        super(PVAM, self).__init__()
-        self.char_num = char_num
-        self.max_length = max_text_length
-        self.num_heads = num_heads
-        self.num_encoder_TUs = num_encoder_tus
-        self.hidden_dims = hidden_dims
-        # Transformer encoder
-        t = 256
-        c = 512
-        self.wrap_encoder_for_feature = WrapEncoderForFeature(
-            src_vocab_size=1,
-            max_length=t,
-            n_layer=self.num_encoder_TUs,
-            n_head=self.num_heads,
-            d_key=int(self.hidden_dims / self.num_heads),
-            d_value=int(self.hidden_dims / self.num_heads),
-            d_model=self.hidden_dims,
-            d_inner_hid=self.hidden_dims,
-            prepostprocess_dropout=0.1,
-            attention_dropout=0.1,
-            relu_dropout=0.1,
-            preprocess_cmd="n",
-            postprocess_cmd="da",
-            weight_sharing=True)
-
-        # PVAM
-        self.flatten0 = paddle.nn.Flatten(start_axis=0, stop_axis=1)
-        self.fc0 = paddle.nn.Linear(
-            in_features=in_channels,
-            out_features=in_channels, )
-        self.emb = paddle.nn.Embedding(
-            num_embeddings=self.max_length, embedding_dim=in_channels)
-        self.flatten1 = paddle.nn.Flatten(start_axis=0, stop_axis=2)
-        self.fc1 = paddle.nn.Linear(
-            in_features=in_channels, out_features=1, bias_attr=False)
-
-    def forward(self, inputs, encoder_word_pos, gsrm_word_pos):
-        b, c, h, w = inputs.shape
-        conv_features = paddle.reshape(inputs, shape=[-1, c, h * w])
-        conv_features = paddle.transpose(conv_features, perm=[0, 2, 1])
-        # transformer encoder
-        b, t, c = conv_features.shape
-
-        enc_inputs = [conv_features, encoder_word_pos, None]
-        word_features = self.wrap_encoder_for_feature(enc_inputs)
-
-        # pvam
-        b, t, c = word_features.shape
-        word_features = self.fc0(word_features)
-        word_features_ = paddle.reshape(word_features, [-1, 1, t, c])
-        word_features_ = paddle.tile(word_features_, [1, self.max_length, 1, 1])
-        word_pos_feature = self.emb(gsrm_word_pos)
-        word_pos_feature_ = paddle.reshape(word_pos_feature,
-                                           [-1, self.max_length, 1, c])
-        word_pos_feature_ = paddle.tile(word_pos_feature_, [1, 1, t, 1])
-        y = word_pos_feature_ + word_features_
-        y = F.tanh(y)
-        attention_weight = self.fc1(y)
-        attention_weight = paddle.reshape(
-            attention_weight, shape=[-1, self.max_length, t])
-        attention_weight = F.softmax(attention_weight, axis=-1)
-        pvam_features = paddle.matmul(attention_weight,
-                                      word_features)  #[b, max_length, c]
-        return pvam_features
-
-
-class GSRM(nn.Layer):
-    def __init__(self, in_channels, char_num, max_text_length, num_heads,
-                 num_encoder_tus, num_decoder_tus, hidden_dims):
-        super(GSRM, self).__init__()
-        self.char_num = char_num
-        self.max_length = max_text_length
-        self.num_heads = num_heads
-        self.num_encoder_TUs = num_encoder_tus
-        self.num_decoder_TUs = num_decoder_tus
-        self.hidden_dims = hidden_dims
-
-        self.fc0 = paddle.nn.Linear(
-            in_features=in_channels, out_features=self.char_num)
-        self.wrap_encoder0 = WrapEncoder(
-            src_vocab_size=self.char_num + 1,
-            max_length=self.max_length,
-            n_layer=self.num_decoder_TUs,
-            n_head=self.num_heads,
-            d_key=int(self.hidden_dims / self.num_heads),
-            d_value=int(self.hidden_dims / self.num_heads),
-            d_model=self.hidden_dims,
-            d_inner_hid=self.hidden_dims,
-            prepostprocess_dropout=0.1,
-            attention_dropout=0.1,
-            relu_dropout=0.1,
-            preprocess_cmd="n",
-            postprocess_cmd="da",
-            weight_sharing=True)
-
-        self.wrap_encoder1 = WrapEncoder(
-            src_vocab_size=self.char_num + 1,
-            max_length=self.max_length,
-            n_layer=self.num_decoder_TUs,
-            n_head=self.num_heads,
-            d_key=int(self.hidden_dims / self.num_heads),
-            d_value=int(self.hidden_dims / self.num_heads),
-            d_model=self.hidden_dims,
-            d_inner_hid=self.hidden_dims,
-            prepostprocess_dropout=0.1,
-            attention_dropout=0.1,
-            relu_dropout=0.1,
-            preprocess_cmd="n",
-            postprocess_cmd="da",
-            weight_sharing=True)
-
-        self.mul = lambda x: paddle.matmul(x=x,
-                                           y=self.wrap_encoder0.prepare_decoder.emb0.weight,
-                                           transpose_y=True)
-
-    def forward(self, inputs, gsrm_word_pos, gsrm_slf_attn_bias1,
-                gsrm_slf_attn_bias2):
-        # ===== GSRM Visual-to-semantic embedding block =====
-        b, t, c = inputs.shape
-        pvam_features = paddle.reshape(inputs, [-1, c])
-        word_out = self.fc0(pvam_features)
-        word_ids = paddle.argmax(F.softmax(word_out), axis=1)
-        word_ids = paddle.reshape(x=word_ids, shape=[-1, t, 1])
-
-        #===== GSRM Semantic reasoning block =====
-        """
-        This module is achieved through bi-transformers,
-        ngram_feature1 is the froward one, ngram_fetaure2 is the backward one
-        """
-        pad_idx = self.char_num
-
-        word1 = paddle.cast(word_ids, "float32")
-        word1 = F.pad(word1, [1, 0], value=1.0 * pad_idx, data_format="NLC")
-        word1 = paddle.cast(word1, "int64")
-        word1 = word1[:, :-1, :]
-        word2 = word_ids
-
-        enc_inputs_1 = [word1, gsrm_word_pos, gsrm_slf_attn_bias1]
-        enc_inputs_2 = [word2, gsrm_word_pos, gsrm_slf_attn_bias2]
-
-        gsrm_feature1 = self.wrap_encoder0(enc_inputs_1)
-        gsrm_feature2 = self.wrap_encoder1(enc_inputs_2)
-
-        gsrm_feature2 = F.pad(gsrm_feature2, [0, 1],
-                              value=0.,
-                              data_format="NLC")
-        gsrm_feature2 = gsrm_feature2[:, 1:, ]
-        gsrm_features = gsrm_feature1 + gsrm_feature2
-
-        gsrm_out = self.mul(gsrm_features)
-
-        b, t, c = gsrm_out.shape
-        gsrm_out = paddle.reshape(gsrm_out, [-1, c])
-
-        return gsrm_features, word_out, gsrm_out
-
-
-class VSFD(nn.Layer):
-    def __init__(self, in_channels=512, pvam_ch=512, char_num=38):
-        super(VSFD, self).__init__()
-        self.char_num = char_num
-        self.fc0 = paddle.nn.Linear(
-            in_features=in_channels * 2, out_features=pvam_ch)
-        self.fc1 = paddle.nn.Linear(
-            in_features=pvam_ch, out_features=self.char_num)
-
-    def forward(self, pvam_feature, gsrm_feature):
-        b, t, c1 = pvam_feature.shape
-        b, t, c2 = gsrm_feature.shape
-        combine_feature_ = paddle.concat([pvam_feature, gsrm_feature], axis=2)
-        img_comb_feature_ = paddle.reshape(
-            combine_feature_, shape=[-1, c1 + c2])
-        img_comb_feature_map = self.fc0(img_comb_feature_)
-        img_comb_feature_map = F.sigmoid(img_comb_feature_map)
-        img_comb_feature_map = paddle.reshape(
-            img_comb_feature_map, shape=[-1, t, c1])
-        combine_feature = img_comb_feature_map * pvam_feature + (
-            1.0 - img_comb_feature_map) * gsrm_feature
-        img_comb_feature = paddle.reshape(combine_feature, shape=[-1, c1])
-
-        out = self.fc1(img_comb_feature)
-        return out
-
-
-class SRNHead(nn.Layer):
-    def __init__(self, in_channels, out_channels, max_text_length, num_heads,
-                 num_encoder_TUs, num_decoder_TUs, hidden_dims, **kwargs):
-        super(SRNHead, self).__init__()
-        self.char_num = out_channels
-        self.max_length = max_text_length
-        self.num_heads = num_heads
-        self.num_encoder_TUs = num_encoder_TUs
-        self.num_decoder_TUs = num_decoder_TUs
-        self.hidden_dims = hidden_dims
-
-        self.pvam = PVAM(
-            in_channels=in_channels,
-            char_num=self.char_num,
-            max_text_length=self.max_length,
-            num_heads=self.num_heads,
-            num_encoder_tus=self.num_encoder_TUs,
-            hidden_dims=self.hidden_dims)
-
-        self.gsrm = GSRM(
-            in_channels=in_channels,
-            char_num=self.char_num,
-            max_text_length=self.max_length,
-            num_heads=self.num_heads,
-            num_encoder_tus=self.num_encoder_TUs,
-            num_decoder_tus=self.num_decoder_TUs,
-            hidden_dims=self.hidden_dims)
-        self.vsfd = VSFD(in_channels=in_channels, char_num=self.char_num)
-
-        self.gsrm.wrap_encoder1.prepare_decoder.emb0 = self.gsrm.wrap_encoder0.prepare_decoder.emb0
-
-    def forward(self, inputs, targets=None):
-        others = targets[-4:]
-        encoder_word_pos = others[0]
-        gsrm_word_pos = others[1]
-        gsrm_slf_attn_bias1 = others[2]
-        gsrm_slf_attn_bias2 = others[3]
-
-        pvam_feature = self.pvam(inputs, encoder_word_pos, gsrm_word_pos)
-
-        gsrm_feature, word_out, gsrm_out = self.gsrm(
-            pvam_feature, gsrm_word_pos, gsrm_slf_attn_bias1,
-            gsrm_slf_attn_bias2)
-
-        final_out = self.vsfd(pvam_feature, gsrm_feature)
-        if not self.training:
-            final_out = F.softmax(final_out, axis=1)
-
-        _, decoded_out = paddle.topk(final_out, k=1)
-
-        predicts = OrderedDict([
-            ('predict', final_out),
-            ('pvam_feature', pvam_feature),
-            ('decoded_out', decoded_out),
-            ('word_out', word_out),
-            ('gsrm_out', gsrm_out),
-        ])
-
-        return predicts