# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang, Wei Kang) # Copyright 2023 Johns Hopkins University (author: Desh Raj) # # See ../../../../LICENSE for clarification regarding multiple authors # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import List, Optional, Tuple import k2 import torch import torch.nn as nn from encoder_interface import EncoderInterface from icefall.utils import add_sos class SURT(nn.Module): """It implements Streaming Unmixing and Recognition Transducer (SURT). https://arxiv.org/abs/2011.13148 """ def __init__( self, mask_encoder: nn.Module, encoder: EncoderInterface, joint_encoder_layer: Optional[nn.Module], decoder: nn.Module, joiner: nn.Module, num_channels: int, encoder_dim: int, decoder_dim: int, joiner_dim: int, vocab_size: int, ): """ Args: mask_encoder: It is the masking network. It generates a mask for each channel of the encoder. These masks are applied to the input features, and then passed to the transcription network. encoder: It is the transcription network in the paper. Its accepts two inputs: `x` of (N, T, encoder_dim) and `x_lens` of shape (N,). It returns two tensors: `logits` of shape (N, T, encoder_dm) and `logit_lens` of shape (N,). decoder: It is the prediction network in the paper. Its input shape is (N, U) and its output shape is (N, U, decoder_dim). It should contain one attribute: `blank_id`. joiner: It has two inputs with shapes: (N, T, encoder_dim) and (N, U, decoder_dim). Its output shape is (N, T, U, vocab_size). Note that its output contains unnormalized probs, i.e., not processed by log-softmax. num_channels: It is the number of channels that the input features will be split into. In general, it should be equal to the maximum number of simultaneously active speakers. For most real scenarios, using 2 channels is sufficient. """ super().__init__() assert isinstance(encoder, EncoderInterface), type(encoder) assert hasattr(decoder, "blank_id") self.mask_encoder = mask_encoder self.encoder = encoder self.joint_encoder_layer = joint_encoder_layer self.decoder = decoder self.joiner = joiner self.num_channels = num_channels self.simple_am_proj = nn.Linear( encoder_dim, vocab_size, ) self.simple_lm_proj = nn.Linear(decoder_dim, vocab_size) self.ctc_output = nn.Sequential( nn.Dropout(p=0.1), nn.Linear(encoder_dim, vocab_size), nn.LogSoftmax(dim=-1), ) def forward_helper( self, x: torch.Tensor, x_lens: torch.Tensor, y: k2.RaggedTensor, prune_range: int = 5, am_scale: float = 0.0, lm_scale: float = 0.0, reduction: str = "sum", beam_size: int = 10, use_double_scores: bool = False, subsampling_factor: int = 1, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Compute transducer loss for one branch of the SURT model. """ encoder_out, x_lens = self.encoder(x, x_lens) assert torch.all(x_lens > 0) if self.joint_encoder_layer is not None: encoder_out = self.joint_encoder_layer(encoder_out) # compute ctc log-probs ctc_output = self.ctc_output(encoder_out) # For the decoder, i.e., the prediction network row_splits = y.shape.row_splits(1) y_lens = row_splits[1:] - row_splits[:-1] blank_id = self.decoder.blank_id sos_y = add_sos(y, sos_id=blank_id) # sos_y_padded: [B, S + 1], start with SOS. sos_y_padded = sos_y.pad(mode="constant", padding_value=blank_id) # decoder_out: [B, S + 1, decoder_dim] decoder_out = self.decoder(sos_y_padded) # Note: y does not start with SOS # y_padded : [B, S] y_padded = y.pad(mode="constant", padding_value=0) y_padded = y_padded.to(torch.int64) boundary = torch.zeros((x.size(0), 4), dtype=torch.int64, device=x.device) boundary[:, 2] = y_lens boundary[:, 3] = x_lens lm = self.simple_lm_proj(decoder_out) am = self.simple_am_proj(encoder_out) with torch.cuda.amp.autocast(enabled=False): simple_loss, (px_grad, py_grad) = k2.rnnt_loss_smoothed( lm=lm.float(), am=am.float(), symbols=y_padded, termination_symbol=blank_id, lm_only_scale=lm_scale, am_only_scale=am_scale, boundary=boundary, reduction=reduction, return_grad=True, ) # ranges : [B, T, prune_range] ranges = k2.get_rnnt_prune_ranges( px_grad=px_grad, py_grad=py_grad, boundary=boundary, s_range=prune_range, ) # am_pruned : [B, T, prune_range, encoder_dim] # lm_pruned : [B, T, prune_range, decoder_dim] am_pruned, lm_pruned = k2.do_rnnt_pruning( am=self.joiner.encoder_proj(encoder_out), lm=self.joiner.decoder_proj(decoder_out), ranges=ranges, ) # logits : [B, T, prune_range, vocab_size] # project_input=False since we applied the decoder's input projections # prior to do_rnnt_pruning (this is an optimization for speed). logits = self.joiner(am_pruned, lm_pruned, project_input=False) with torch.cuda.amp.autocast(enabled=False): pruned_loss = k2.rnnt_loss_pruned( logits=logits.float(), symbols=y_padded, ranges=ranges, termination_symbol=blank_id, boundary=boundary, reduction=reduction, ) # Compute ctc loss supervision_segments = torch.stack( ( torch.arange(len(x_lens), device="cpu"), torch.zeros_like(x_lens, device="cpu"), torch.clone(x_lens).detach().cpu(), ), dim=1, ).to(torch.int32) # We need to sort supervision_segments in decreasing order of num_frames indices = torch.argsort(supervision_segments[:, 2], descending=True) supervision_segments = supervision_segments[indices] # Works with a BPE model decoding_graph = k2.ctc_graph(y, modified=False, device=x.device) dense_fsa_vec = k2.DenseFsaVec( ctc_output, supervision_segments, allow_truncate=subsampling_factor - 1, ) ctc_loss = k2.ctc_loss( decoding_graph=decoding_graph, dense_fsa_vec=dense_fsa_vec, output_beam=beam_size, reduction="none", use_double_scores=use_double_scores, ) return (simple_loss, pruned_loss, ctc_loss) def forward( self, x: torch.Tensor, x_lens: torch.Tensor, y: k2.RaggedTensor, prune_range: int = 5, am_scale: float = 0.0, lm_scale: float = 0.0, reduction: str = "sum", beam_size: int = 10, use_double_scores: bool = False, subsampling_factor: int = 1, return_masks: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Args: x: A 3-D tensor of shape (N, T, C). x_lens: A 1-D tensor of shape (N,). It contains the number of frames in `x` before padding. y: A ragged tensor of shape (N*num_channels, S). It contains the labels of the N utterances. The labels are in the range [0, vocab_size). All the channels are concatenated together one after another. prune_range: The prune range for rnnt loss, it means how many symbols(context) we are considering for each frame to compute the loss. am_scale: The scale to smooth the loss with am (output of encoder network) part lm_scale: The scale to smooth the loss with lm (output of predictor network) part reduction: "sum" to sum the losses over all utterances in the batch. "none" to return the loss in a 1-D tensor for each utterance in the batch. beam_size: The beam size used in CTC decoding. use_double_scores: If True, use double precision for CTC decoding. subsampling_factor: The subsampling factor of the model. It is used to compute the supervision segments for CTC loss. return_masks: If True, return the masks as well as masked features. Returns: Return the transducer loss. Note: Regarding am_scale & lm_scale, it will make the loss-function one of the form: lm_scale * lm_probs + am_scale * am_probs + (1-lm_scale-am_scale) * combined_probs """ assert x.ndim == 3, x.shape assert x_lens.ndim == 1, x_lens.shape assert y.num_axes == 2, y.num_axes assert x.size(0) == x_lens.size(0), (x.size(), x_lens.size()) # Apply the mask encoder B, T, F = x.shape processed = self.mask_encoder(x) # B,T,F*num_channels masks = processed.view(B, T, F, self.num_channels).unbind(dim=-1) x_masked = [x * m for m in masks] # Recognition # Stack the inputs along the batch axis h = torch.cat(x_masked, dim=0) h_lens = torch.cat([x_lens for _ in range(self.num_channels)], dim=0) simple_loss, pruned_loss, ctc_loss = self.forward_helper( h, h_lens, y, prune_range, am_scale, lm_scale, reduction=reduction, beam_size=beam_size, use_double_scores=use_double_scores, subsampling_factor=subsampling_factor, ) # Chunks the outputs into 2 parts along batch axis and then stack them along a new axis. simple_loss = torch.stack( torch.chunk(simple_loss, self.num_channels, dim=0), dim=0 ) pruned_loss = torch.stack( torch.chunk(pruned_loss, self.num_channels, dim=0), dim=0 ) ctc_loss = torch.stack(torch.chunk(ctc_loss, self.num_channels, dim=0), dim=0) if return_masks: return (simple_loss, pruned_loss, ctc_loss, x_masked, masks) else: return (simple_loss, pruned_loss, ctc_loss, x_masked)