Retrieving submodule from module declared with partial

[rolandgvc]

Given the following example:

class Net(nn.Module):
  ...
  
  def setup(self):
      self.backbone = Backbone(...)
      self.classifier = nn.Dense(...)

    def __call__(self, inputs):
      x = self.backbone(inputs)
      x = self.classifier(x)
      return x

def create_model(rng, private_params, **kwargs):
  model = functools.partial(Net, private_params, **kwargs)
  init_batch = jnp.ones((1, 224, 224, 3), jnp.float32)
  params = Net(**kwargs).init(rng, init_batch)
  return model, params

def net_50(rng, **kwargs):
  return _create_model(rng, private_params, **kwargs)

model, params = net_50(RNG, train=False)

How would it be possible to extract the backbone submodule from model while still using functools.partial to hide the private params from the user?

[avital]

Hi Roland, can you return an instance of Net instead a function that creates it? We need to document this better but just calling Dense(3) doesn't actually create any variables, it's just a skeleton "module template" (what we call an "unbound module" in the code)?

Though I think at the moment you still won't have access to Net().backbone because setup is only called when the module is bound... This actually may be good to improve.

Can you say more about your full use-case for accessing net.backbone in your code?

[avital]

(I tried doing !pip install -q git+https://github.com/rolandgvc/flaxvision/tree/linen/ in a Colab but then I got: ERROR: flaxvision 0.0.5 has requirement flax==0.2.0, but you'll have flax 0.3.0 which is incompatible.)

[rolandgvc]

Right, sorry for that. I'm putting a colab together now. Thanks for taking the time!

[rolandgvc]

Here it is https://colab.research.google.com/drive/1KAQspBG3BEJ2PgNNtd1n63kgf78PPpZU?usp=sharing

[avital]

Thanks Roland! I revised your example and got it working (I had to change a few other things as well -- you didn't need @compact at all around SegmentationModel.__call__, and your call to init should have model.init(rng, init_batch) ): https://colab.research.google.com/drive/1QlwImBjxZxMFkM50qR9FzJ4o-QW51Waa#scrollTo=J4So4-Sf8oUN

The important part of the change is this pattern:

class SegmentationModel(nn.Module):
    backbone_fn: Any
    classifier_fn: Any

    def setup(self):
        self.backbone = self.backbone_fn()
        self.classifier = self.classifier_fn()

    def __call__(self, inputs, train: bool = False):
        x = self.backbone(inputs, train=train)
        x = self.classifier(x, train=train)
        return x

def _make_model(rng, train, **kwargs):
    resnet_model, _ = resnet50(rng)
    backbone_fn = lambda: ResNet.make_backbone(resnet_model)
    classifier_fn = lambda: FCNHead(10, train=train)

    return SegmentationModel(backbone_fn, classifier_fn, **kwargs)

That is -- pass in a function that makes a module as an attribute into a larger module, and then assign it during setup. The assignment during setup is what attaches it to the parent module, so there must be explicit assignment. Simply passing in attributes doesn't go through that (intentionally - because you may use a submodule somewhere else in the module tree and then it's incorrect to bind it twice).

We really need to improve this user experience, and I hope to spend some more time on #686, thanks for this motivating example!

Please let me know if this last suggestion seems to work for you. If you hit any more issues, please continue to post here (or as a new discussion, as appropriate).

Thanks for slogging with us -- this is extremely useful feedback and we need this to keep improving Flax!

[rolandgvc]

Awesome!! Thanks @avital and the team for always being active and helpful in these discussions! I will try the example in a bit. In case I still struggle with it, I will create a new discussion and link it here (since this one deviated a bit from the original topic). Many thanks!