Function visualize_attn

utils/vis.py:115–157

Visualize cross attention: `stage_id`th downsampling block, mean over all timesteps starting from step start, `block_id`th Transformer block, second item (conditioned), mean over heads, show each token cross_attention_probs_tensors: One of `cross_attention_probs_down_tensors`, `cross_at

(
    token_map,
    cross_attention_probs_tensors,
    stage_id,
    block_id,
    visualize_step_start=10,
    input_ca_has_condition_only=False,
)

# This function has not been adapted to new `saved_attn`.
def visualize_attn(
    token_map,
    cross_attention_probs_tensors,
    stage_id,
    block_id,
    visualize_step_start=10,
    input_ca_has_condition_only=False,
):
    """
    Visualize cross attention: `stage_id`th downsampling block, mean over all timesteps starting from step start, `block_id`th Transformer block, second item (conditioned), mean over heads, show each token
    cross_attention_probs_tensors:
    One of `cross_attention_probs_down_tensors`, `cross_attention_probs_mid_tensors`, and `cross_attention_probs_up_tensors`
    stage_id: index of downsampling/mid/upsaming block
    block_id: index of the transformer block
    """

    plt.figure(figsize=(20, 8))

    for token_id in range(len(token_map)):
        token = token_map[token_id]
        plt.subplot(1, len(token_map), token_id + 1)
        plt.title(token)
        attn = cross_attention_probs_tensors[stage_id][visualize_step_start:].mean(
            dim=0
        )[block_id]

        if not input_ca_has_condition_only:
            assert (
                attn.shape[0] == 2
            ), f"Expect to have 2 items (uncond and cond), but found {attn.shape[0]} items"
            attn = attn[1]
        else:
            assert (
                attn.shape[0] == 1
            ), f"Expect to have 1 item (cond only), but found {attn.shape[0]} items"
            attn = attn[0]

        attn = attn.mean(dim=0)[:, token_id]
        H = W = int(math.sqrt(attn.shape[0]))
        attn = attn.reshape((H, W))
        plt.imshow(attn.cpu().numpy())

    plt.show()


# This function has not been adapted to new `saved_attn`.