Note

You are reading the documentation for MMOCR 0.x, which will soon be deprecated by the end of 2022. We recommend you upgrade to MMOCR 1.0 to enjoy fruitful new features and better performance brought by OpenMMLab 2.0. Check out the maintenance plan, changelog, code and documentation of MMOCR 1.0 for more details.

Source code for mmocr.models.textrecog.layers.robust_scanner_fusion_layer

# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
from mmcv.runner import BaseModule


class RobustScannerFusionLayer(BaseModule):

    def __init__(self, dim_model, dim=-1, init_cfg=None):
        super().__init__(init_cfg=init_cfg)

        self.dim_model = dim_model
        self.dim = dim

        self.linear_layer = nn.Linear(dim_model * 2, dim_model * 2)
        self.glu_layer = nn.GLU(dim=dim)

    def forward(self, x0, x1):
        assert x0.size() == x1.size()
        fusion_input = torch.cat([x0, x1], self.dim)
        output = self.linear_layer(fusion_input)
        output = self.glu_layer(output)

        return output