Class Transformer

src/shape_assembly/models/train/transformer.py:146–173

        return self.w_2(self.norm(F.relu(self.w_1(x)).transpose(2, 1).contiguous()).transpose(2, 1).contiguous())

class Transformer(pl.LightningModule):

    def __init__(self, cfg):
        super().__init__()
        c = copy.deepcopy
        attn = MultiHeadedAttention(
            cfg.model.num_heads, 
            cfg.model.pc_feat_dim
        )

        ff = PositionwiseFeedForward(
            cfg.model.pc_feat_dim, 
            cfg.model.transformer_feat_dim
        )

        self.model = EncoderDecoder(
            Encoder(EncoderLayer(cfg.model.pc_feat_dim, c(attn), c(ff)), cfg.model.num_blocks),
            Decoder(DecoderLayer(cfg.model.pc_feat_dim, c(attn), c(attn), c(ff)), cfg.model.num_blocks),
            nn.Sequential(),                        
            nn.Sequential(),                        
            nn.Sequential()
        )

    def forward(self, src, tgt):
        src = src.transpose(2, 1).contiguous()
        tgt = tgt.transpose(2, 1).contiguous()
        src_corr_feat = self.model(tgt, src, None, None).transpose(2, 1).contiguous()
        return src_corr_feat