Access and replace layers

[gianlucadetommaso]

Hello,
given an arbitrary Flax model, is it possible to access, and potentially replace, the layers it formed of?

For example, consider the following simple model:

import flax.linen as nn

class Model(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(features=10)(x)
        x = nn.relu(x)
        x = nn.Dense(features=1)(x)
        return x
        
model = Model()     

Given model, is there any method that would allow to access the sequence of layers used (nn.Dense(features=10) and nn.Dense(features=1)), and potentially replace them?

Thanks! 🙏

[chiamp]

Hi @gianlucadetommaso, Module.bind can allow you to access the layers, but I don't know of a way to replace them. Hope this helps!

[levskaya]

By convention we define layers by code execution. Other groups have written simple template systems on top of flax to allow fine-grained config using advanced config systems like "gin" or "fiddle". But we don't offer that out of the box.

[gianlucadetommaso]

Thanks for letting me know.

[gianlucadetommaso]

Please bear with me and let me iterate on this once more. I would like to be sure there is no way around this issue.

Consider a scenario where a user would like to inspect a model and add a nn.Dropout layer whenever a nn.Dense layer is encountered. Can you think about any simple way to do this in Flax?

[cgarciae]

@gianlucadetommaso Flax's lazy behavior makes it such that you cannot do this in general. You can only replace things in when they are inputs to a module (i.e. dataclass fields):

import flax.linen as nn
import jax
import jax.numpy as jnp

class Foo(nn.Module):
  bar: nn.Module

  def __call__(self, x):
    return self.bar(x)

x = jnp.ones((1, 32, 32, 3))

module = Foo(bar=nn.Dense(10))
module.bar = nn.Conv(features=10, kernel_size=(3, 3), padding='SAME') # replace
variables = module.init(jax.random.PRNGKey(0), x)

print("\nFoo\n-----------")
print(jax.tree_map(lambda x: x.shape, variables))

#--------------------
# Sequential
#--------------------
module = nn.Sequential([
    nn.Dense(10),
    nn.relu,
    nn.Dense(20),
])
module.layers = [
    nn.Conv(features=10, kernel_size=(3, 3), padding='SAME') if isinstance(x, nn.Dense) else x
    for x in module.layers
]

variables = module.init(jax.random.PRNGKey(0), x)
print("\nSequential\n-----------")
print(jax.tree_map(lambda x: x.shape, variables))

However, you cannot easily (if at all) replace modules that are defined inside setup or @compact as they are only materialized inside init / apply.

[gianlucadetommaso]

Thanks for the exhaustive answer!

[gianlucadetommaso]

@cgarciae As a follow up question, let's assume I can give up on working with an instantiated model, but rather I can work directly with the model class (e.g. Model from my initial example). In such case, is there any good way to inspect Model, find some particular layers and replace them?

For example, similarly to your example above, I may want to find all Dense layers and replace them.

[cgarciae]

@gianlucadetommaso unless you provide an API for it (e.g. having a class argument or passing the modules as inputs), currently there is no way to replace submodules inside created inside compact or setup.