Method forward

cdslib/core/nn/modules/attention.py:95–220

Args: inputs: (seq_len, batch_size, dim_input) input used to compute attention context_vectors: (num_chars, batch_size, dim_context) can be onehot encoding or embeddings init_means: (batch_size, num_mixtures) or (1, batch_s

(
            self,
            inputs: torch.Tensor,
            context_vectors: torch.Tensor,
            init_means: torch.Tensor = None,
            extra_chars: int = 0,
    )

        return torch.zeros(batch_size, self.num_mixtures, device=device)

    def forward(
            self,
            inputs: torch.Tensor,
            context_vectors: torch.Tensor,
            init_means: torch.Tensor = None,
            extra_chars: int = 0,
    ):
        """
        Args:
            inputs: (seq_len, batch_size, dim_input)
                input used to compute attention
            context_vectors: (num_chars, batch_size, dim_context)
                can be onehot encoding or embeddings
            init_means: (batch_size, num_mixtures) or (1, batch_size, num_mixtures)
                current means of the gaussians (None: all zeros)
            extra_chars:
                how many extra char_idx to calculate

        Returns:
            attn_contexts (seq_len, batch_size, dim_context)
                one for each time step
            attn_weights (seq_len, batch_size, total_char)
                one for each time step, total_char = num_char + extra_chars
            means (seq_len, batch_size, num_mixtures)
                mean of each gaussian window
            vars (seq_len, batch_size, num_mixtures)
                variance of each gaussian window
            weights (seq_len, batch_size, num_mixtures)
                weights of each gaussian window
        """
        seq_len = inputs.size(0)
        batch_size = inputs.size(1)
        num_char = context_vectors.size(0)
        dim_context = context_vectors.size(2)

        total_char = num_char + extra_chars

        attn_weight_dict = self.compute_attn_weights(
            inputs=inputs,
            total_char=total_char,
            init_means=init_means,
        )

        attn_contexts = self.compute_attn_contexts(
            attn_weights=attn_weight_dict['attn_weights'],
            context_vectors=context_vectors,
        )

        return (
            attn_contexts,
            attn_weight_dict['attn_weights'],
            attn_weight_dict['means'],
            attn_weight_dict['vars'],
            attn_weight_dict['weights'],
        )

        # if init_means is None:
        #     init_means = self.get_init_means(batch_size, device=inputs.device)  # (batch_size,num_mixtures)