mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-26 18:24:18 +00:00
Replace torchaudio rnnt_loss to k2 pruned rnnt loss
This commit is contained in:
parent
5b6699a835
commit
2d6631ac76
@ -54,7 +54,8 @@ def greedy_search(model: Transducer, encoder_out: torch.Tensor) -> List[int]:
|
|||||||
# fmt: off
|
# fmt: off
|
||||||
current_encoder_out = encoder_out[:, t:t+1, :]
|
current_encoder_out = encoder_out[:, t:t+1, :]
|
||||||
# fmt: on
|
# fmt: on
|
||||||
logits = model.joiner(current_encoder_out, decoder_out)
|
logits = model.joiner(
|
||||||
|
current_encoder_out.unsqueeze(2), decoder_out.unsqueeze(1))
|
||||||
# logits is (1, 1, 1, vocab_size)
|
# logits is (1, 1, 1, vocab_size)
|
||||||
|
|
||||||
log_prob = logits.log_softmax(dim=-1)
|
log_prob = logits.log_softmax(dim=-1)
|
||||||
|
@ -30,21 +30,17 @@ class Joiner(nn.Module):
|
|||||||
"""
|
"""
|
||||||
Args:
|
Args:
|
||||||
encoder_out:
|
encoder_out:
|
||||||
Output from the encoder. Its shape is (N, T, C).
|
Output from the encoder. Its shape is (N, T, S_range, C) for
|
||||||
|
training and (1, 1, 1, C) for decoding.
|
||||||
decoder_out:
|
decoder_out:
|
||||||
Output from the decoder. Its shape is (N, U, C).
|
Output from the decoder. Its shape is (N, T, S_range, C) for
|
||||||
|
training and (1, 1, 1, C) for decoding.
|
||||||
Returns:
|
Returns:
|
||||||
Return a tensor of shape (N, T, U, C).
|
Return a tensor of shape (N, T, S_range, C) for training and
|
||||||
|
(1, 1, 1, C) for decoding.
|
||||||
"""
|
"""
|
||||||
assert encoder_out.ndim == decoder_out.ndim == 3
|
assert encoder_out.ndim == decoder_out.ndim == 4
|
||||||
assert encoder_out.size(0) == decoder_out.size(0)
|
assert encoder_out.shape == decoder_out.shape
|
||||||
assert encoder_out.size(2) == decoder_out.size(2)
|
|
||||||
|
|
||||||
encoder_out = encoder_out.unsqueeze(2)
|
|
||||||
# Now encoder_out is (N, T, 1, C)
|
|
||||||
|
|
||||||
decoder_out = decoder_out.unsqueeze(1)
|
|
||||||
# Now decoder_out is (N, 1, U, C)
|
|
||||||
|
|
||||||
logit = encoder_out + decoder_out
|
logit = encoder_out + decoder_out
|
||||||
logit = torch.tanh(logit)
|
logit = torch.tanh(logit)
|
||||||
|
@ -14,15 +14,9 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# 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 k2
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
import torchaudio
|
|
||||||
import torchaudio.functional
|
|
||||||
from encoder_interface import EncoderInterface
|
from encoder_interface import EncoderInterface
|
||||||
|
|
||||||
from icefall.utils import add_sos
|
from icefall.utils import add_sos
|
||||||
@ -102,24 +96,34 @@ class Transducer(nn.Module):
|
|||||||
|
|
||||||
decoder_out, _ = self.decoder(sos_y_padded)
|
decoder_out, _ = self.decoder(sos_y_padded)
|
||||||
|
|
||||||
logits = self.joiner(encoder_out, decoder_out)
|
|
||||||
|
|
||||||
# rnnt_loss requires 0 padded targets
|
# rnnt_loss requires 0 padded targets
|
||||||
# Note: y does not start with SOS
|
# Note: y does not start with SOS
|
||||||
y_padded = y.pad(mode="constant", padding_value=0)
|
y_padded = y.pad(mode="constant", padding_value=0)
|
||||||
|
|
||||||
assert hasattr(torchaudio.functional, "rnnt_loss"), (
|
boundary = torch.zeros((x.size(0), 4),
|
||||||
f"Current torchaudio version: {torchaudio.__version__}\n"
|
dtype=torch.int64, device=x.device)
|
||||||
"Please install a version >= 0.10.0"
|
boundary[:, 2] = y_lens
|
||||||
)
|
boundary[:, 3] = x_lens
|
||||||
|
|
||||||
loss = torchaudio.functional.rnnt_loss(
|
y_padded = y_padded.to(torch.int64)
|
||||||
logits=logits,
|
|
||||||
targets=y_padded,
|
simple_loss, (px_grad, py_grad) = k2.rnnt_loss_simple(
|
||||||
logit_lengths=x_lens,
|
decoder_out, encoder_out, y_padded, blank_id, boundary, True)
|
||||||
target_lengths=y_lens,
|
|
||||||
blank=blank_id,
|
# TODO: make s_range configurable
|
||||||
reduction="sum",
|
ranges = k2.get_rnnt_prune_ranges(px_grad, py_grad, x_lens, s_range=5)
|
||||||
)
|
|
||||||
|
am_pruned, lm_pruned = k2.do_rnnt_pruning(
|
||||||
|
encoder_out, decoder_out, ranges)
|
||||||
|
|
||||||
|
logits = self.joiner(am_pruned, lm_pruned)
|
||||||
|
|
||||||
|
# boundary may change after pruning
|
||||||
|
boundary[:, 2] = ranges[:,-1,-1]
|
||||||
|
|
||||||
|
pruned_loss = k2.rnnt_loss_pruned(
|
||||||
|
logits, y_padded, ranges, blank_id, boundary)
|
||||||
|
|
||||||
|
loss = -(torch.sum(simple_loss) + torch.sum(pruned_loss))
|
||||||
|
|
||||||
return loss
|
return loss
|
||||||
|
Loading…
x
Reference in New Issue
Block a user