« Back to API reference

neuraltrain.models.common.BahdanauAttention

class neuraltrain.models.common.BahdanauAttention(input_size, hidden_size, use_queries: bool = False)

Bahdanau attention from [1].

Implementation inspired from pytorch’s seq2seq tutorial: https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html#the-decoder

[1]

Bahdanau, Dzmitry, Kyunghyun Cho, and Yoshua Bengio. “Neural machine translation by jointly learning to align and translate.” arXiv preprint arXiv:1409.0473 (2014).

forward(keys, queries=None)

Parameters:

• keys – Key tensor of shape (batch_size, n_features, n_times).

• queries – Optional query tensor of shape (batch_size, n_features, n_times). If None, only keys are used.

« Back to API reference