# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang) # # 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. """ Note we use `rnnt_loss` from torchaudio, which exists only in torchaudio >= v0.10.0. It also means you have to use torch >= v1.10.0 """ import k2 import torch import torch.nn as nn import torchaudio import torchaudio.functional from encoder_interface import EncoderInterface from icefall.utils import add_sos class Transducer(nn.Module): """It implements https://arxiv.org/pdf/1211.3711.pdf "Sequence Transduction with Recurrent Neural Networks" """ def __init__( self, encoder: EncoderInterface, decoder: nn.Module, joiner: nn.Module, ): """ Args: encoder: It is the transcription network in the paper. Its accepts two inputs: `x` of (N, T, C) and `x_lens` of shape (N,). It returns two tensors: `logits` of shape (N, T, C) 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, C). It should contain two attributes: `blank_id` and `sos_id`. joiner: It has two inputs with shapes: (N, T, C) and (N, U, C). Its output shape is (N, T, U, C). Note that its output contains unnormalized probs, i.e., not processed by log-softmax. """ super().__init__() assert isinstance(encoder, EncoderInterface) assert hasattr(decoder, "blank_id") assert hasattr(decoder, "sos_id") self.encoder = encoder self.decoder = decoder self.joiner = joiner def forward( self, x: torch.Tensor, x_lens: torch.Tensor, y: k2.RaggedTensor, ) -> 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 with 2 axes [utt][label]. It contains labels of each utterance. Returns: Return the transducer loss. """ 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) == y.dim0 encoder_out, x_lens = self.encoder(x, x_lens) assert torch.all(x_lens > 0) # Now 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_id = self.decoder.sos_id sos_y = add_sos(y, sos_id=sos_id) sos_y_padded = sos_y.pad(mode="constant", padding_value=blank_id) sos_y_padded = sos_y_padded.to(torch.int64) decoder_out, _ = self.decoder(sos_y_padded) logits = self.joiner(encoder_out, decoder_out) # rnnt_loss requires 0 padded targets # Note: y does not start with SOS y_padded = y.pad(mode="constant", padding_value=0) assert hasattr(torchaudio.functional, "rnnt_loss"), ( f"Current torchaudio version: {torchaudio.__version__}\n" "Please install a version >= 0.10.0" ) loss = torchaudio.functional.rnnt_loss( logits=logits, targets=y_padded, logit_lengths=x_lens, target_lengths=y_lens, blank=blank_id, reduction="sum", ) return loss