Grab Attention Maps from Nested PyTorch Modules

What it does Visualizer is a small PyTorch utility that yanks attention maps—or any local variable—out of deeply nested model definitions without making you thread return values through layers or register forward hooks on every block. You decorate the function that builds the map, activate the tracer before importing your model, and the variable lands in a cache after the forward pass.

The interesting bit Instead of wrestling with module paths like model.blocks.3.attn or polluting your code with global variables, it manipulates Python bytecode to intercept local variables at runtime. The decorator is effectively inert during training—if you don’t activate it before import, it leaves no trace and imposes zero overhead.

Key highlights

• Pulls variables by name from inside functions, regardless of how deeply nested the module hierarchy is.
• Captures all attention maps across Transformer layers in one shot by decorating the shared forward method once.
• Zero training-time footprint: leave the decorator in place; without activate() it does nothing.
• Not limited to attention—any local variable in any function is fair game.

Caveats

• The variable you want must keep its name until the end of the function; if another statement reassigns to the same name, you get the final value, not the intermediate one.
• You must call get_local.activate() before importing the model, because the decorator binds at import time.

Verdict Worth a look if you’re debugging or visualizing attention in Vision Transformers and tired of hook bookkeeping. Skip it if you already have a clean hook pipeline or don’t need to peek inside specific local variables.