Linen api, siamese networks and transfer learning

[mortner31]

Hello, I don't know if this is the right place where to put these comments, anyhow, maybe it may be of some use for your team.

First, I am having a lot of fun playing with the new linen interface of Flax. Congrats to the team for your work, this is huge.
Everything is cristal clear, it is good to have the ability to know exactly what is going on, step by steps.
As a side note, I often train networks for production purposes, and I am really anxious on how big frameworks like Keras and Torch hide things underneath. I really need to interact with the code at many levels in order to control what is going on, both at training, inference, debugging and unit testing, hence my interest in your work.

From the architecture point of view, I am interested in two big things:

• transfer learning from siamese network with weak losses (triplet or contrastive)
• switching architecture while training (as done in once-for-all https://arxiv.org/abs/1908.09791) for which I believe flax will be of great use, since the parameter structure is really simple.

Following many discussions, I have been able to implement a toy example quite easily for the first case ie. siamese network instanciation and transfert. Here is the example : https://colab.research.google.com/drive/1rplADiwAoDI8byx_foHE-Eoq2e_6ur4y?usp=sharing

While playing with the example, it raised a couple of questions:

1. Can a module __call__ return multiple output such as a tuple x1,x2 or a dict of outputs , are there any constraints or best practice ? Should it be avoided ?

2. While transfering weights from a network to another I was able to overwrite the module definition (n_filters of Conv) without raising any error. I think it might be in flax philosophy to have sanity checks while initializing module with params, in order to avoid silently overwriting the module definition with init parameters...

3. I have a mixed feeling about the fact that the keys of the nested dict of params can represent two different things : the name of a block of parameters in lower levels of the tree or the type of parameters in the first nodes (for instance batch_stats vs params). I love the fact that it is simple to handle, but in the end , I feel like I will have to write a couple of utility functions to handle the structure (such as split in N different trees, apply a function, and merge back - I actually had to write something like this in the transfert example). Maybe there is something to figure out here or maybe the current state is so easy to understand that it is the best way for doing things.

[cgarciae]

I think this discussion is relevant:

#437

[mortner31]

Yes, I actually mentioned your interesting discussion in my colab notebook, since I took the transfer example from your discussion.
I did not know if I should start a new discussion, or put this post at the bottom of yours ... If you think it is better practice to add it at the bottom of yours, I can delete this one and move it at the bottom of yours.

[avital]

Hi @totor31 -- thanks for reaching out, and sorry for the late reply!

1. Can a module __call__ return multiple output such as a tuple x1,x2 or a dict of outputs , are there any constraints or best practice ? Should it be avoided ?

It's up to you -- I've seen projects where the full model returns logics, endpoints -- where endpoints is a dict containing the output of various helpful intermediate layers. Flax/Linen has no opinion about this.

2. While transfering weights from a network to another I was able to overwrite the module definition (n_filters of Conv) without raising any error. I think it might be in flax philosophy to have sanity checks while initializing module with params, in order to avoid silently overwriting the module definition with init parameters...

I agree -- do you mind filing a GitHub issue so that we can track this? IIUC you're asking for an assertion that the shapes of variables are the same as what's defined in the call to self.param, if they are already present in the variable dict

3. I have a mixed feeling about the fact that the keys of the nested dict of params can represent two different things : the name of a block of parameters in lower levels of the tree or the type of parameters in the first nodes (for instance batch_stats vs params). I love the fact that it is simple to handle, but in the end , I feel like I will have to write a couple of utility functions to handle the structure (such as split in N different trees, apply a function, and merge back - I actually had to write something like this in the transfert example). Maybe there is something to figure out here or maybe the current state is so easy to understand that it is the best way for doing things.

Indeed it's a bit annoying to do this kind of "top-level variable dict surgery" when manipulation params and other state separately. We chatted about this a bit but couldn't find a better pattern. If you (or anyone else) has any suggestions please do let us know! [Edit: we may have a new optimizer pattern that helps with this]

[mortner31]

thx for the answer, I will fill an issue

[mortner31]

here for the record #477