We want to visualize internal model calculations, e.g. attention maps, in a research project with multiple different models. The goal is to create a common framework for these visualizations with the following constraints:

they should only be calculated and stored during validation to save computation time during training
the interface should be common for all models (i.e. PyTorch modules with a forward function)
it should be possible to automatically visualize the generated images in tensorboard
it should work (e.g. correct naming of generated images) with modules calling other submodules

I am wondering what is the best way to implement this (i.e. the most usable and generic way). Basically I had the idea of passing a dictionary to the module into which the module can write, depending on it’s current mode.
Did you already encounter such a problem? If yes how did you implement this?
Thanks!
Jan