How to List Full Model Topology, including Activation Layers

[oelghaza]

I am looking for a way to reconstruct my NN model while using the new NNX API. If I have a basic MLP with Activations between Linear layers, I am not able to retrieve any information on the Activation layers.

I understand this is due to them being JAX functions rather than NNX Modules. Is there a way for me to perform model surgery and retrieve information about the activation layers of a saved model — without explicitly creating a Flax Module for each activation layer?

Ie if my model has the following architecture:

-Linear1
-Activation1
-Linear2
-Activation2

Using the model surgery techniques in the guides(display(), iter_children(), iter_modules(), saving/restoring checkpoint) I can only retrieve the following:

-Linear1
-Linear2

[cgarciae]

Hey! Flax NNX will print all instance attributes of a Module so if you refactor your code such that the activation functions are attributes you should be able to find them with nnx.iter_graph.

[oelghaza]

I tried this and I am able to see the attributes listed in the order they were created. This does not show me the layer order for a forward pass of the network.

I am looking to view the actual layer-by-layer topology of the neural network.

[cgarciae]

Can you show an example of what is going wrong?

[oelghaza]

Hi, I just picked this back up. Thanks for helping!
I attached some screenshots to give a better idea.

In this example I have a basic MLP that alternates between Dense and Activation Layers. All the model surgery techniques I have tried list all the dense, followed by all the activation layers. I am not able to view the actual topology which interleaves the activation layers between the dense ones; I also cannot find the activation as an attribute of the desne layers.

I.e. I want to take in a model and see the topology of a forward pass: [dense1, relu1, dense2, relu2, ...] instead of what I see [dense1, dense2, relu1, relu2].

Model def:

iter_graph() output:

graphdef output:

[pablordoricaw]

Hey @oelghaza, here's a way to implement a neural network to see the anatomy of a forward pass when you look at/traverse the model.

@cgarciae - I'd appreciate your thoughts on this approach.

The idea is to leverage the nnx.Sequential class in combination with your relu_module class, and an additional softmax_module class:

from flax import nnx

class relu_module(nnx.Module):
    def __call__(self, x):
        return nnx.relu(x)

class softmax_module(nnx.Module):
    def __call__(self, x):
        return nnx.softmax(x)
        
nnx_1_jet_mlp = nnx.Sequential(
    nnx.Linear(64, 32, rngs=nnx.Rngs(0)),
    relu_module(),

    nnx.Linear(64, 32, rngs=nnx.Rngs(0)),
    relu_module(),

    nnx.Linear(64, 32, rngs=nnx.Rngs(0)),
    relu_module(),

    nnx.Linear(64, 32, rngs=nnx.Rngs(0)),
    softmax_module()
)

If you create your model as an instance of the nnx.Sequential class, then the GraphDef for the nnx_1_jet_mlp model has the layers sorted in the order of the forward pass:

Traversing it with the iter_modules() method: