Recommended way to partially load weights into a model ?

Hi,

Is there a recommended way to partially load weights to a new model ?

At the moment I have an inept solution which uses the fact that fc attribute is the only eqx.nn.Linear instance in my model

is_leaf = lambda x : True if isinstance(x, eqx.nn.Linear) else False
new_model = tree_deserialise_leaves('model.eqx', CNN(num_classes=26, key=key), is_leaf=is_leaf)

This approach won’t work where there are multiple Linear layers. Is there a recommended way to perform such loading?

A name parameter for modules might be useful here.

[patrick-kidger]

Yep, I could see some documentation being a good idea. (Including a version of my above example?)

[patrick-kidger]

Right – so the answer is that whilst tree_map just-so-happens to be ordered, this is actually an implementation detail. That is, JAX reserves the right to change this behaviour about ordering in the future.

In contrast tree_flatten does guarantee that it will return things in order, and JAX actually includes a test that this is the case. This is from here: https://github.com/google/jax/pull/11658 (Which as you can see is very recent. The context is that I raised this question internally with the JAX team, and it was agreed to guarantee that tree_flatten will preserve ordering, but not to offer the same guarantee for tree_map .)

The current implementation of tree_map is a lot neater than my approach though! I suggest we copy-paste that for our ordered_tree_map, instead of the one I wrote.