Emformer¶
- class torchaudio.models.Emformer(input_dim: int, num_heads: int, ffn_dim: int, num_layers: int, segment_length: int, dropout: float = 0.0, activation: str = 'relu', left_context_length: int = 0, right_context_length: int = 0, max_memory_size: int = 0, weight_init_scale_strategy: Optional[str] = 'depthwise', tanh_on_mem: bool = False, negative_inf: float = -100000000.0)[source]¶
Emformer 架构,首次提出于 Emformer: Efficient Memory Transformer Based Acoustic Model for Low Latency Streaming Speech Recognition [Shi et al., 2021]。
另请参阅
torchaudio.pipelines.RNNTBundle
: 包含预训练模型的 ASR Pipeline。
- 参数:
input_dim (int) – 输入维度。
num_heads (int) – 每个 Emformer 层的注意力头数量。
ffn_dim (int) – 每个 Emformer 层的 Feedforward 网络隐藏层维度。
num_layers (int) – 要实例化的 Emformer 层数量。
segment_length (int) – 每个输入段的长度。
dropout (float, optional) – dropout 概率。(默认值:0.0)
activation (str, optional) – 每个 Emformer 层的 Feedforward 网络中使用的激活函数。必须是 (“relu”, “gelu”, “silu”) 之一。(默认值:“relu”)
left_context_length (int, optional) – 左上下文长度。(默认值:0)
right_context_length (int, optional) – 右上下文长度。(默认值:0)
max_memory_size (int, optional) – 要使用的最大内存元素数量。(默认值:0)
weight_init_scale_strategy (str or None, optional) – 每层权重初始化缩放策略。必须是 (“depthwise”, “constant”,
None
) 之一。(默认值:“depthwise”)tanh_on_mem (bool, optional) – 如果
True
,则对内存元素应用 tanh。(默认值:False
)negative_inf (float, optional) – 在注意力权重中用于负无穷大的值。(默认值:-1e8)
示例
>>> emformer = Emformer(512, 8, 2048, 20, 4, right_context_length=1) >>> input = torch.rand(128, 400, 512) # batch, num_frames, feature_dim >>> lengths = torch.randint(1, 200, (128,)) # batch >>> output, lengths = emformer(input, lengths) >>> input = torch.rand(128, 5, 512) >>> lengths = torch.ones(128) * 5 >>> output, lengths, states = emformer.infer(input, lengths, None)
方法¶
forward¶
- Emformer.forward(input: Tensor, lengths: Tensor) Tuple[Tensor, Tensor] ¶
用于训练和非流式推理的前向传播。
B: 批次大小;T: 批次中最大输入帧数;D: 每帧的特征维度。
- 参数:
input (torch.Tensor) – 经过右上下文帧右填充的 utterance 帧,形状为 (B, T + right_context_length, D)。
lengths (torch.Tensor) – 形状为 (B,),其中第 i 个元素表示
input
中第 i 个批次元素的有效 utterance 帧数量。
- 返回值:
- Tensor
输出帧,形状为 (B, T, D)。
- Tensor
输出长度,形状为 (B,),其中第 i 个元素表示输出帧中第 i 个批次元素的有效帧数量。
- 返回类型:
(Tensor, Tensor)
infer¶
- Emformer.infer(input: Tensor, lengths: Tensor, states: Optional[List[List[Tensor]]] = None) Tuple[Tensor, Tensor, List[List[Tensor]]] ¶
用于流式推理的前向传播。
B: 批次大小;D: 每帧的特征维度。
- 参数:
input (torch.Tensor) – 经过右上下文帧右填充的 utterance 帧,形状为 (B, segment_length + right_context_length, D)。
lengths (torch.Tensor) – 形状为 (B,),其中第 i 个元素表示
input
中第 i 个批次元素的有效帧数量。states (List[List[torch.Tensor]] or None, optional) – 表示在之前调用
infer
中生成的内部状态的张量列表的列表。(默认值:None
)
- 返回值:
- Tensor
输出帧,形状为 (B, segment_length, D)。
- Tensor
输出长度,形状为 (B,),其中第 i 个元素表示输出帧中第 i 个批次元素的有效帧数量。
- List[List[Tensor]]
输出状态;表示在当前调用
infer
中生成的内部状态的张量列表的列表。
- 返回类型:
(Tensor, Tensor, List[List[Tensor]])