mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-08 09:32:20 +00:00
1d44da845b
* Begin to add RNN-T training for librispeech. * Copy files from conformer_ctc. Will edit it. * Use conformer/transformer model as encoder. * Begin to add training script. * Add training code. * Remove long utterances to avoid OOM when a large max_duraiton is used. * Begin to add decoding script. * Add decoding script. * Minor fixes. * Add beam search. * Use LSTM layers for the encoder. Need more tunings. * Use stateless decoder. * Minor fixes to make it ready for merge. * Fix README. * Update RESULT.md to include RNN-T Conformer. * Minor fixes. * Fix tests. * Minor fixes. * Minor fixes. * Fix tests.
128 lines
4.0 KiB
Python
128 lines
4.0 KiB
Python
# 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
|
|
|
|
assert hasattr(torchaudio.functional, "rnnt_loss"), (
|
|
f"Current torchaudio version: {torchaudio.__version__}\n"
|
|
"Please install a version >= 0.10.0"
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
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)
|
|
|
|
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
|