espnet2.asr.layers.fastformer.FastSelfAttention

source

class espnet2.asr.layers.fastformer.FastSelfAttention(size, attention_heads, dropout_rate)

Bases: Module

Fast self-attention used in Fastformer.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

espnet_initialization_fn()

forward(xs_pad, mask)

Forward method.

• Parameters:
  • xs_pad – (batch, time, size = n_heads * attn_dim)
  • mask – (batch, 1, time), nonpadding is 1, padding is 0
• Returns: (batch, time, size)
• Return type: torch.Tensor

init_weights(module)

transpose_for_scores(x)

Reshape and transpose to compute scores.

• Parameters:x – (batch, time, size = n_heads * attn_dim)
• Returns: (batch, n_heads, time, attn_dim)