Class ACTORDecoder

diffplanner/models/transformers/actor.py:173–245

@SUBMODULES.register_module()
class ACTORDecoder(BaseModule):

    def __init__(self,
                 max_seq_len=16,
                 njoints=None,
                 nfeats=None,
                 input_feats=None,
                 input_dim=256,
                 latent_dim=256,
                 condition_dim=None,
                 num_heads=4,
                 ff_size=1024,
                 num_layers=8,
                 activation='gelu',
                 dropout=0.1,
                 use_condition=False,
                 num_class=None,
                 pos_embedding='sinusoidal',
                 init_cfg=None):
        super().__init__(init_cfg=init_cfg)
        if input_dim != latent_dim:
            self.linear = nn.Linear(input_dim, latent_dim)
        else:
            self.linear = nn.Identity()
        self.njoints = njoints
        self.nfeats = nfeats
        if input_feats is None:
            assert self.njoints is not None and self.nfeats is not None
            self.input_feats = njoints * nfeats
        else:
            self.input_feats = input_feats
        self.max_seq_len = max_seq_len
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.condition_dim = condition_dim
        self.use_condition = use_condition
        self.num_class = num_class
        if self.use_condition:
            if num_class is None:
                self.condition_bias = build_MLP(condition_dim, latent_dim)
            else:
                self.condition_bias = nn.Parameter(torch.randn(num_class, latent_dim))
        if pos_embedding == 'sinusoidal':
            self.pos_encoder = SinusoidalPositionalEncoding(latent_dim, dropout)
        else:
            self.pos_encoder = LearnedPositionalEncoding(latent_dim, dropout, max_len=max_seq_len)
        seqTransDecoderLayer = nn.TransformerDecoderLayer(
            d_model=self.latent_dim,
            nhead=num_heads,
            dim_feedforward=ff_size,
            dropout=dropout,
            activation=activation)
        self.seqTransDecoder = nn.TransformerDecoder(
            seqTransDecoderLayer,
            num_layers=num_layers)

        self.final = nn.Linear(self.latent_dim, self.input_feats)