Tight couplings in the design of network modules

[activatedgeek]

Thank you for the great work on this library.

I want to open up a discussion about a few design decisions of tight couplings with the nn.Module definitions, and suggest alternatives.

In its current form, the model definition is tighly coupled with a few moving parts:

1. Normalization resnet.py#L297: This option should not be allowed since it can lead to unknown side-effects. It should be the responsibility of the data loading routines to handle the data standardization to provide clean separation of concerns.
2. Pretraining Class Selection resnet.py#L302: Here again, the model creation code should be independent of external considerations. Nothing should be overridden, at least in the model definition. I think it is completely ok to have separate library functions which implement this behavior outside the nn.Module.
3. Parameter Initialization resnet.py#L314: Parameters are explicitly decoupled in how FLAX processes a model.apply({ 'params': params, **other_vars}, ....) call. Similarly, if we can process the checkpoints into FrozenDict's that can be loaded via flax.training.checkpoints. This makes things like transfer learning easier because I can directly modify the params PyTree in the FrozenDict instead of worrying about the size of the output from the nn.Module. i.e. keep everything except the last layer params the same and just reinit the last layer params using an appropriate initializer.

What do you think?

[matthias-wright]

Hey @activatedgeek, thanks for reaching out!

1. I added this option because I always find it annoying to look up the mean and std for a given dataset. If a user wants to handle the normalization in their data pipeline, they can simply set normalize=False.
2. The idea was that if pretrained is set to 'imagenet', then num_classes has to be 1000. I would be okay with removing this and only logging a warning instead.
3. Yes, I agree.

Cheers!