Method _input_stats_summaries

axlearn/audio/decoder_asr.py:239–271

Computes input lengths stats. Args: input_batch: See forward method signature. target_paddings: See _compute_target_paddings method return value. is_valid_example: A 0/1 Tensor of shape [batch_size], 1 if the example is a valid input to th

(
        self, input_batch: Nested[Tensor], *, target_paddings: Tensor, is_valid_example: Tensor
    )

        raise NotImplementedError(type(self))

    def _input_stats_summaries(
        self, input_batch: Nested[Tensor], *, target_paddings: Tensor, is_valid_example: Tensor
    ) -> dict[str, Union[WeightedSummary, Tensor]]:
        """Computes input lengths stats.

        Args:
            input_batch: See forward method signature.
            target_paddings: See _compute_target_paddings method return value.
            is_valid_example: A 0/1 Tensor of shape [batch_size], 1 if the example is
                a valid input to the loss computation.

        Returns:
            A dictionary of input stats summaries.
        """
        valid_frames = (1.0 - input_batch["paddings"]) * is_valid_example[:, None]
        valid_labels = (1.0 - target_paddings) * is_valid_example[:, None]

        total_source_lengths = jnp.sum(valid_frames)
        total_target_lengths = jnp.sum(valid_labels)
        total_num_examples = jnp.maximum(is_valid_example.sum(), 1.0)
        total_num_frames = jnp.maximum(jnp.size(input_batch["paddings"]), 1)
        input_stats = {
            "input_stats/average_target_length": WeightedSummary(
                total_target_lengths / total_num_examples, total_num_examples
            ),
            "input_stats/average_source_length": WeightedSummary(
                total_source_lengths / total_num_examples, total_num_examples
            ),
            "input_stats/frame_packing_efficiency": WeightedSummary(
                total_source_lengths / total_num_frames, total_num_frames
            ),
        }
        return input_stats


class CTCDecoderModel(BaseASRDecoderModel):