Function _add_codebook_summaries

axlearn/common/quantizer.py:276–310

Helper function to compute codebook distribution statistics and add to summaries. The statistics are from all frames, not only on those masked frames in self-supervised training. # ToDo(zhiyunlu): Add support to take an additional mask input. Args: context: Module invocation co

(*, context: InvocationContext, onehots: Tensor, paddings: Tensor)

def _add_codebook_summaries(*, context: InvocationContext, onehots: Tensor, paddings: Tensor):
    """Helper function to compute codebook distribution statistics and add to summaries.

    The statistics are from all frames, not only on those masked frames in self-supervised training.
    # ToDo(zhiyunlu): Add support to take an additional mask input.

    Args:
        context: Module invocation context to add summaries to.
        onehots: onehot of BaseQuantizer.Output.ids.
        paddings: 0/1 tensor of shape [batch_size, seq_len], where 0 is valid position.
    """
    coverage = compute_code_coverage(onehots=onehots, paddings=paddings)
    pplx, entropy = compute_code_pplx(onehots=onehots, paddings=paddings)
    batch_size = paddings.shape[0]

    num_frames = jnp.sum(safe_not(paddings))
    context.add_summary(
        "codebook/num_frames",
        WeightedSummary(num_frames.astype(jnp.float32) / batch_size, batch_size),
    )
    # Mean coverage of all codebooks.
    context.add_summary(
        "codebook/coverage",
        WeightedSummary(coverage, jnp.maximum(1, num_frames)),
    )
    # Mean perplexity of all codebooks.
    context.add_summary(
        "codebook/pplx",
        WeightedSummary(pplx, jnp.maximum(1, num_frames)),
    )
    # Mean entropy of all codebooks.
    context.add_summary(
        "codebook/entropy",
        WeightedSummary(entropy, jnp.maximum(1, num_frames)),
    )


class RandomVectorQuantizer(BaseQuantizer):