fastNLP.modules.attention module

undocumented

class fastNLP.modules.attention.MultiHeadAttention(d_model: int = 512, n_head: int = 8, dropout: float = 0.0, layer_idx: int = None)[源代码]

基类:torch.nn.modules.module.Module

Attention is all you need中提到的多头注意力

forward(query, key, value, key_mask=None, attn_mask=None, state=None)[源代码]
参数
  • query – batch x seq x dim

  • key – batch x seq x dim

  • value – batch x seq x dim

  • key_mask – batch x seq 用于指示哪些key不要attend到;注意到mask为1的地方是要attend到的

  • attn_mask – seq x seq, 用于mask掉attention map。 主要是用在训练时decoder端的self attention,下三角为1

  • state – 过去的信息,在inference的时候会用到,比如encoder output、decoder的prev kv。这样可以减少计算。

返回

training: bool
class fastNLP.modules.attention.BiAttention[源代码]

基类:torch.nn.modules.module.Module

Bi Attention module

对于给定的两个向量序列 \(a_i\)\(b_j\) , BiAttention模块将通过以下的公式来计算attention结果

\[\begin{split}\begin{array}{ll} \\ e_{ij} = {a}^{\mathrm{T}}_{i}{b}_{j} \\ {\hat{a}}_{i} = \sum_{j=1}^{\mathcal{l}_{b}}{\frac{\mathrm{exp}(e_{ij})}{\sum_{k=1}^{\mathcal{l}_{b}}{\mathrm{exp}(e_{ik})}}}{b}_{j} \\ {\hat{b}}_{j} = \sum_{i=1}^{\mathcal{l}_{a}}{\frac{\mathrm{exp}(e_{ij})}{\sum_{k=1}^{\mathcal{l}_{a}}{\mathrm{exp}(e_{ik})}}}{a}_{i} \\ \end{array}\end{split}\]
forward(premise_batch, premise_mask, hypothesis_batch, hypothesis_mask)[源代码]
参数
  • premise_batch (torch.Tensor) – [batch_size, a_seq_len, hidden_size]

  • premise_mask (torch.Tensor) – [batch_size, a_seq_len]

  • hypothesis_batch (torch.Tensor) – [batch_size, b_seq_len, hidden_size]

  • hypothesis_mask (torch.Tensor) – [batch_size, b_seq_len]

返回

torch.Tensor attended_premises: [batch_size, a_seq_len, hidden_size] torch.Tensor attended_hypotheses: [batch_size, b_seq_len, hidden_size]

training: bool
class fastNLP.modules.attention.SelfAttention(input_size, attention_unit=300, attention_hops=10, drop=0.5, initial_method=None)[源代码]

基类:torch.nn.modules.module.Module

这是一个基于论文 A structured self-attentive sentence embedding 的Self Attention Module.

__init__(input_size, attention_unit=300, attention_hops=10, drop=0.5, initial_method=None)[源代码]
参数
  • input_size (int) – 输入tensor的hidden维度

  • attention_unit (int) – 输出tensor的hidden维度

  • attention_hops (int) –

  • drop (float) – dropout概率,默认值为0.5

  • initial_method (str) – 初始化参数方法

forward(input, input_origin)[源代码]
参数
  • input (torch.Tensor) – [batch_size, seq_len, hidden_size] 要做attention的矩阵

  • input_origin (torch.Tensor) – [batch_size, seq_len] 原始token的index组成的矩阵,含有pad部分内容

Return torch.Tensor output1

[batch_size, multi-head, hidden_size] 经过attention操作后输入矩阵的结果

Return torch.Tensor output2

[1] attention惩罚项,是一个标量

training: bool