espnet2.asr.encoder.e_branchformer_encoder.EBranchformerEncoderLayer

source

class espnet2.asr.encoder.e_branchformer_encoder.EBranchformerEncoderLayer(size: int, attn: Module, cgmlp: Module, feed_forward: Module | None, feed_forward_macaron: Module | None, dropout_rate: float, merge_conv_kernel: int = 3)

Bases: Module

E-Branchformer encoder layer module.

• Parameters:
  • size (int) – model dimension
  • attn – standard self-attention or efficient attention
  • cgmlp – ConvolutionalGatingMLP
  • feed_forward – feed-forward module, optional
  • feed_forward – macaron-style feed-forward module, optional
  • dropout_rate (float) – dropout probability
  • merge_conv_kernel (int) – kernel size of the depth-wise conv in merge module

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

forward(x_input, mask, cache=None)

Compute encoded features.

• Parameters:
  • x_input (Union *[*Tuple , torch.Tensor ]) – Input tensor w/ or w/o pos emb.
    • w/ pos emb: Tuple of tensors [(#batch, time, size), (1, time, size)].
    • w/o pos emb: Tensor (#batch, time, size).
  • mask (torch.Tensor) – Mask tensor for the input (#batch, 1, time).
  • cache (torch.Tensor) – Cache tensor of the input (#batch, time - 1, size).
• Returns: Output tensor (#batch, time, size). torch.Tensor: Mask tensor (#batch, time).
• Return type: torch.Tensor