Method __init__

openrec/modeling/decoders/abinet_decoder.py:157–203

(self,
                 in_channels,
                 out_channels,
                 nhead=8,
                 num_layers=3,
                 dim_feedforward=2048,
                 dropout=0.1,
                 max_length=25,
                 iter_size=3,
                 **kwargs)

class ABINetDecoder(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 nhead=8,
                 num_layers=3,
                 dim_feedforward=2048,
                 dropout=0.1,
                 max_length=25,
                 iter_size=3,
                 **kwargs):
        super().__init__()
        self.max_length = max_length + 1
        d_model = in_channels
        self.pos_encoder = PositionalEncoding(dropout=0.1, dim=d_model)
        self.encoder = nn.ModuleList([
            TransformerBlock(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=dim_feedforward,
                attention_dropout_rate=dropout,
                residual_dropout_rate=dropout,
                with_self_attn=True,
                with_cross_attn=False,
            ) for _ in range(num_layers)
        ])
        self.decoder = PositionAttention(
            max_length=self.max_length,  # additional stop token
            in_channels=d_model,
            num_channels=d_model // 8,
            mode='nearest',
        )
        self.out_channels = out_channels
        self.cls = nn.Linear(d_model, self.out_channels)
        self.iter_size = iter_size
        if iter_size > 0:
            self.language = BCNLanguage(
                d_model=d_model,
                nhead=nhead,
                num_layers=4,
                dim_feedforward=dim_feedforward,
                dropout=dropout,
                max_length=max_length,
                num_classes=self.out_channels,
            )
            # alignment
            self.w_att_align = nn.Linear(2 * d_model, d_model)
            self.cls_align = nn.Linear(d_model, self.out_channels)

    def forward(self, x, data=None):
        # bs, c, h, w