Visualizing modules

[TheRisenPhoenix]

Hi! :)
I was wondering what the best way to visualize (e.g. summary or onnx export/vis) torchrl modules is.
I tried both, torchinfo and onnx export on my policy module (which is a MultiAgentConvNet and -MLP), but it fails. Maybe I'm doing something wrong, but possibly that's just not the right way to do it? Any advice on this? :)

Edit:
The doc says that just using print(cnn) would return a nice overview of the layers, but when I run the example, I get a nested view of tensordicts?

from torchrl.modules import MultiAgentConvNet
n_agents = 2
cnn = MultiAgentConvNet(
     n_agents,
     centralised = True,
     share_params = True
)
print(cnn)

[vmoens]

Hello!
Happy to chat about this, I don't think we carefully considered it yet.
I agree currently this specific module just prints the params - which is a tensordict. We should have a better way to print that.
I guess that because the module is sort of functional (we pull out the params to be able to vmap over them) we would need to patch the util you're using or tell it what to print.

Can I ask what's the use case? For instance "With a regular module such as X, I can do Y, but not with this one" (to know what's out objective and deliverable here)

[TheRisenPhoenix]

Thanks for your readiness to help!
For me, the use case is to have a functionality to visualize (or at least print) the layers of an architecture.
Say I'd like to use some fancy architecture for the neural network powering my policy, built out of CNNs and MLPs. Usually, I can summarize the architecture (and the amount of parameters of each layer) for usual torch models using a simple print(model) statement, and that's also what is promised by the documentation (down at examples). However, there are also other packages such as torchinfo which can do this. For more advanced visualization, you can use the onnx export and load it into other tools to get some graphical representation.

I understand that the onnx export might be the most difficult (although also the most beautiful one), but perhaps some kind of print summary would be helpful, to see of which layers (with how many parameters) a model (module) consists of?
Actually what the examples states would already be sufficient (at least for me :D ):

>>> print(cnn)
MultiAgentConvNet(
    (agent_networks): ModuleList(
        (0): ConvNet(
        (0): Conv2d(4, 32, kernel_size=(5, 5), stride=(2, 2))
        (1): ELU(alpha=1.0)
        (2): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
        (3): ELU(alpha=1.0)
        (4): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
        (5): ELU(alpha=1.0)
        (6): SquashDims()
        )
    )
)

[vmoens]

Does this help?
#2192