visTvis

Store and visualize attention weights from Vision Transformers in a model-agnostic way. The package ships two decorators:

• @visTvis_store saves a return value tensor to disk with a predictable folder layout.
• @visTvis_layer_counter injects an int-like counter that keeps track of which layer is currently active so saved files end up in the right slot.
• Visualization helpers overlay attention on input frames and stack results into PDFs.

Installation

pip install vistvis

For local development in this repo:

make install

Quick start

Decorate the attention routine you want to log and point it to a base directory:

from vistvis import visTvis_store


@visTvis_store(base_folder_path="./runs", identifier="demo", file_name="attn.pt", folder_name="attn")
def attn_func(attn_weight):
    # ... compute attn_weight ...
    return attn_weight

Next, wrap the function that calls your attention blocks so it receives a LayerCounter. The counter behaves like an int (supports indexing, comparisons, += 1, etc.) and updates the visTvis layer state on every mutation:

from vistvis import visTvis_layer_counter


@visTvis_layer_counter(identifier="demo", var_name="layer_idx")
def run_blocks(blocks, inputs, layer_idx=0):
    for block in blocks:
        attn_weight = block(inputs)
        attn_func(attn_weight)  # stored under runs/demo/attn/layer_<idx>/attn.pt
        layer_idx += 1  # updates the visTvis layer state
    return int(layer_idx)

During runtime each call to attn_func will save attn_weight under:

<base_folder_path>/<identifier>/<folder_name>/layer_<layer_idx>/<file_name>

By default the layer counter starts at zero and persists between calls for the same identifier. If you need manual control you can use from vistvis import state and call state.update(identifier, value) before invoking decorated functions.

Visualize saved attention

Save attention under attn/ and inputs under input/ so the helpers can discover them:

from vistvis import AttentionPath, run_overlay_for_reconstruction, plot_attention_for_reconstruction, StaticMetadata

paths = AttentionPath(base_dir="./runs", experiment_name="demo")
metadata = StaticMetadata(patch_size=16, num_special_tokens=0)  # or set env VIS_TVIS_METADATA_PATH to a YAML file

run_overlay_for_reconstruction(
    paths,
    query_frame_idx=0,
    query_idx=0,  # or "interactive" to select via GUI
    key_frame_idx=0,
    head="mean",  # or int / "all"
    output_path="./overlays_demo",
    metadata=metadata,
)

plot_attention_for_reconstruction("./overlays_demo")

Demo

Choose Query Patches Interactively from frame N

Get per-layer-head attention overlays on frame M

Publishing

• Update pyproject.toml metadata.
• Build and upload with python -m build followed by twine upload dist/*.