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)

Emformer 架構在 Emformer: Efficient Memory Transformer Based Acoustic Model for Low Latency Streaming Speech Recognition ([Shi et al.，2021]) 中介紹。

另請參閱

• emformer_rnnt_model()、emformer_rnnt_base()：工廠函式。

• torchaudio.pipelines.RNNTBundle：具有預訓練模型的 ASR pipeline。

參數:

• input_dim (int) – 輸入維度。

• num_heads (int) – 每個 Emformer 層中的注意力頭數量。

• ffn_dim (int) – 每個 Emformer 層的前饋網路的隱藏層維度。

• num_layers (int) – 要實例化的 Emformer 層數量。

• segment_length (int) – 每個輸入片段的長度。

• dropout (float, optional) – dropout 機率。（預設值：0.0）

• activation (str, optional) – 在每個 Emformer 層的前饋網路中使用的激活函數。必須是 (“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 或 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]

用於訓練和非串流推論的 Forward pass。

B：批次大小；T：批次中輸入影格的最大數量；D：每個影格的特徵維度。

參數:

• input (torch.Tensor) – 以右側上下文影格右側填充的語音影格，形狀為 (B, T + right_context_length, D)。

• lengths (torch.Tensor) – 形狀為 (B,)，第 i 個元素表示 input 中第 i 個批次元素的有效語音影格數。

返回:

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]]]

用於串流推論的 Forward pass。

B：批次大小；D：每個影格的特徵維度。

參數:

• input (torch.Tensor) – 以右側上下文影格右側填充的語音影格，形狀為 (B, segment_length + right_context_length, D)。

• lengths (torch.Tensor) – 形狀為 (B,)，第 i 個元素表示 input 中第 i 個批次元素的有效影格數。

• states (List[List[torch.Tensor]] or None, optional) – 代表先前調用 infer 中產生的內部狀態的 tensors 列表。（預設值：None）

返回:

Tensor

輸出影格，形狀為 (B, segment_length, D)。

Tensor

輸出長度，形狀為 (B,)，第 i 個元素表示輸出影格中第 i 個批次元素的有效影格數。

List[List[Tensor]]

輸出狀態；代表目前調用 infer 中產生的內部狀態的 tensors 列表。

返回類型:

(Tensor, Tensor, List[List[Tensor]])