Method forward

axlearn/audio/decoder_asr.py:832–905

Computes the transducer loss. Args: input_batch: A dict containing: inputs: A Tensor of shape [batch_size, num_frames, dim]. paddings: A 0/1 Tensor of shape [batch_size, num_frames]. 1's represent paddings. target_labels: An int Te

(self, input_batch: Nested[Tensor])

        self._add_child("transducer", transducer_cfg)

    def forward(self, input_batch: Nested[Tensor]) -> tuple[Tensor, Nested[Tensor]]:
        """Computes the transducer loss.

        Args:
            input_batch: A dict containing:
                inputs: A Tensor of shape [batch_size, num_frames, dim].
                paddings: A 0/1 Tensor of shape [batch_size, num_frames]. 1's represent paddings.
                target_labels: An int Tensor of shape [batch_size, num_labels]. Prediction target
                    of the transducer decoder.
                target: A dictionary with input_ids as key, and an int Tensor of shape
                    [batch_size, num_labels] as value. Prediction inputs to the transducer decoder.

            For both target_labels and target["input_ids"], values should be in the range
            [0, vocab_size). target_labels does not contain BOS and valid label tokens are
            followed by a EOS token. input_ids starts with a BOS token. Sequences are not
            truncated. Out-of-range values are excluded from the loss calculation.

        Returns:
            A tuple (loss, per_example):
                loss: A scalar of the transducer loss.
                per_example: A dict containing transducer decoder outputs of the following keys:
                    weight: A tensor of shape [batch_size], the aggregation weight of the
                        per-example loss.
                    loss: A tensor of shape [batch_size] representing per-example loss.
        """
        cfg: TransducerDecoderModel.Config = self.config

        # [batch, src_max_len, joint_dim].
        am_data = self.am_proj(input_batch["inputs"])
        am_paddings: Tensor = input_batch["paddings"]
        chex.assert_type(am_paddings, jnp.bool)
        target_labels: Tensor = input_batch["target_labels"]
        target_paddings: Tensor = _compute_target_paddings(target_labels, vocab_size=cfg.vocab_size)

        # [batch, tgt_max_len, joint_dim].
        lm_data = self.lm_proj(self.prediction_network(inputs=input_batch["target"]["input_ids"]))

        _, per_example = self.transducer(
            am_data=am_data,
            am_paddings=am_paddings,
            lm_data=lm_data,
            lm_paddings=target_paddings,
            target_labels=target_labels,
        )
        per_example_loss, per_example_weight = (
            per_example["loss"],
            per_example["is_valid_example"],
        )
        per_example_weight = per_example_weight.astype(per_example_loss.dtype)

        # Compute weighted loss.
        loss = jnp.sum(per_example_loss * per_example_weight) / jnp.maximum(
            per_example_weight.sum(), 1
        )
        aux_outputs = dict(per_example_weight=per_example_weight, per_example_loss=per_example_loss)

        # Add input summaries.
        input_summary = self._input_stats_summaries(