`trainable` for BatchNorm stops parameters from being saved and loaded

[xukai92]

Below is defined to only take gradient of the β and γ in batch norm layers.

trainable(bn::BatchNorm) = (bn.β, bn.γ)

However, this stops us from using params and loadparams! to save and load parameters as the other two fields, μ and σ², which are updated during training as well, to be saved and loaded.

Maybe it's just fine to not define trainable(bn::BatchNorm) = (bn.β, bn.γ) as μ and σ² doesn't seems to have gradient?

[CarloLucibello]

Maybe is fine to add μ and σ² to the trainable parameters, since they are used only at test time and they have zero gradient during training.
A more general solution would be to add a flag to params to also return non-trainable fields.

[xukai92]

Maybe is fine to add μ and σ² to the trainable parameters, since they are used only at test time and they have zero gradient during training.

Yes that's what I mean ... don't know why I typed "no" even I said "doesn't` in the end of the setence...

[CarloLucibello]

@xukai92 what's your specific use case? can you work around it with bson saving and loading?

[xukai92]

I used to have issues using BSON.jl - I will try it again.

[nantonel]

Apparently, if you save your model weights:

julia> weights = params(model);

julia> using BSON: @save

julia> @save "mymodel.bson" weights

the batchnorm parameters are lost but if you save with @save "mymodel.bson" model they are preserved.

[ToucheSir]

We should probably advocate against saving implicit Params in favour of saving explicit params or the model itself. Flat params are just too brittle and are easily invalidated with architectural changes.

[darsnack]

I suppose we should think more comprehensively about how to do this. It has come up in FluxML/MetalheadWeights#2. On one hand, Metalhead should just save the model itself and return that instead copying over the weights. On the other hand, this makes it possible to silently return a different model with pretrain = true or pretrain = false. Hopefully, we would never make a release that allows that to happen, but it has the code smell of a poor design.

Ideally, the saving mechanism should allow users to specify state that isn't trainable to be saved. And the loading mechanism should error when the architectures don't match. e.g. maybe we should be doing a structural walk of the model instead of a flat iteration over the parameters.

[ToucheSir]

That all makes sense to me. I think it (mostly the structural walk) might even be possible to do with what's currently in Optimisers.jl.

Ideally, the saving mechanism should allow users to specify state that isn't trainable to be saved. And the loading mechanism should error when the architectures don't match.

PyTorch handles this by returning a set of mismatched locations when loading and optionally throwing when it detects a mismatch. I think we're on the same page about wanting to do something similar for Flux?

Ideally, the saving mechanism should allow users to specify state that isn't trainable to be saved.

This is where it gets "fun" IMO. Trainable state is probably (can't think of a counterexample) a subset of saveable state. For example norm layer statistics aren't trainable, but you'd definitely want to persist them. The simplest solution would be to use the (fields...,) part of @functor StructType (fields...) more so that fmapstructure gives you back only those params which are worth saving. Depending on the serialization format, I think a further pass would still be required to filter out types like functions.

[darsnack]

For loading, I think using fmap(f, m1, m2) with the latest changes on Functors should throw an error when the structures don't match, right?

For saving more than trainable, isn't the easiest to just define savable 😅? It can default to trainable (I agree that "trainable" is a subset of "savable"). Maybe we can modify params (of course this breaks implicit params...I meant have some saveparams(m, path)) to use fmapstructure with the appropriate "subset walk".

[ToucheSir]

I'm not sure there's a need for a new function, functor can probably fill that role on its own :)

[darsnack]

True, but functor should be the most expansive set of fields and not all of them are worth saving. Maybe this distinction is not so important though. It's certainly safer to save all of them.