# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang, Wei Kang) # # 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. import k2 import torch from torch import Tensor import torch.nn as nn from encoder_interface import EncoderInterface from scaling import ScaledLinear from diagonalize import get_diag_covar_in, apply_transformation_in, get_transformation, apply_transformation_in, apply_transformation_out 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, encoder_dim: int, decoder_dim: int, joiner_dim: int, vocab_size: int, ): """ Args: 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. """ super().__init__() assert isinstance(encoder, EncoderInterface), type(encoder) assert hasattr(decoder, "blank_id") self.encoder = encoder self.decoder = decoder self.joiner = joiner self.simple_am_proj = ScaledLinear( encoder_dim, vocab_size, initial_speed=0.5 ) self.simple_lm_proj = ScaledLinear(decoder_dim, vocab_size) 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, warmup: float = 1.0, ) -> 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. 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 warmup: A value warmup >= 0 that determines which modules are active, values warmup > 1 "are fully warmed up" and all modules will be active. 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) == y.dim0 encoder_out, x_lens = self.encoder(x, x_lens, warmup=warmup) 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_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="sum", 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="sum", ) return (simple_loss, pruned_loss) def diagonalize(self) -> None: cur_transform = None for l in self.encoder.encoder.layers: if cur_transform is not None: l.apply_transformation_in(cur_transform) cur_transform = l.get_transformation_out() l.apply_transformation_out(cur_transform) self.encoder.diagonalize() # diagonalizes self_attn layers, this is # purely internal to the self_attn layers. apply_transformation_in(self.simple_am_proj, cur_transform) apply_transformation_in(self.joiner.encoder_proj, cur_transform) def _test_model(): import logging logging.getLogger().setLevel(logging.INFO) from conformer import Conformer from joiner import Joiner from decoder import Decoder feature_dim = 40 attention_dim = 256 encoder_dim = 512 decoder_dim = 513 joiner_dim = 514 vocab_size = 1000 encoder = Conformer(num_features=40, subsampling_factor=4, d_model=encoder_dim, nhead=4, dim_feedforward=512, num_encoder_layers=4) decoder = Decoder( vocab_size=600, decoder_dim=decoder_dim, blank_id=0, context_size=2) joiner = Joiner( encoder_dim=encoder_dim, decoder_dim=decoder_dim, joiner_dim=joiner_dim, vocab_size=vocab_size) model = Transducer(encoder=encoder, decoder=decoder, joiner=joiner, encoder_dim=encoder_dim, decoder_dim=decoder_dim, joiner_dim=joiner_dim, vocab_size=vocab_size) batch_size = 5 seq_len = 50 feats = torch.randn(batch_size, seq_len, feature_dim) x_lens = torch.full((batch_size,), seq_len, dtype=torch.int64) y = k2.ragged.create_ragged_tensor(torch.arange(5, dtype=torch.int32).reshape(1,5).expand(batch_size,5)) model.eval() # eval mode so it's not random. (simple_loss1, pruned_loss1) = model(feats, x_lens, y) model.diagonalize() (simple_loss2, pruned_loss2) = model(feats, x_lens, y) print(f"simple_loss1 = {simple_loss1.mean().item()}, simple_loss2 = {simple_loss2.mean().item()}") print(f"pruned_loss1 = {pruned_loss1.mean().item()}, pruned_loss2 = {pruned_loss2.mean().item()}") model.diagonalize() if __name__ == '__main__': _test_model()