Method forward

axlearn/audio/decoder_asr.py:353–417

Computes CTC 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 sha

(
        self,
        input_batch: Nested[Tensor],
    )

        return ret_dict

    def forward(
        self,
        input_batch: Nested[Tensor],
    ) -> tuple[Tensor, Nested[Tensor]]:
        """Computes CTC 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].
                    Values should be in the range [0, vocab_size). We assume there are no BOS
                    tokens, and that sequences are not truncated. Out-of-range values are excluded
                    from the loss calculation (e.g., paddings and EOS can be represented this way).

        Returns:
            A tuple (loss, aux_outputs):
                loss: A scalar loss value.
                aux_outputs: A dict containing:
                    per_example_loss: A float Tensor of shape [batch_size].
                    per_example_weight: A float Tensor of shape [batch_size].
        """
        cfg: CTCDecoderModel.Config = self.config
        paddings: Tensor = input_batch["paddings"]
        target_labels: Tensor = input_batch["target_labels"]
        target_paddings: Tensor = _compute_target_paddings(target_labels, vocab_size=cfg.vocab_size)

        # Compute CTC loss.
        logits = self.predict(input_batch)
        per_example_loss = optax.ctc_loss(
            logits=logits,
            logit_paddings=paddings,
            labels=target_labels,
            label_paddings=target_paddings,
            blank_id=cfg.blank_id,
        )

        # Drop examples with targets longer than inputs.
        per_example_weight = _is_valid_ctc_seq(
            paddings=paddings, target_labels=target_labels, target_paddings=target_paddings
        )
        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 summaries.
        summary = self._input_stats_summaries(
            input_batch=input_batch,
            target_paddings=target_paddings,
            is_valid_example=per_example_weight,
        )
        summary.update(
            self._loss_summaries(
                total_ctc_loss=jnp.sum(per_example_loss * per_example_weight),
                per_example_weight=per_example_weight,