Method __init__

src/shape_assembly/models/train/transformer.py:148–167

(self, cfg)

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()