Make U matrices not persistent to reduce `state_dict` size. about mace HOT 8 OPEN

Hey @pfebrer,

That is a reasonable solution. Could you try it and check if it solves your issues?

from mace.

It indeed solved the problem, my checkpoint files size has been reduced from 680 MB to 26 MB :) I could also restart the training without problems.

But maybe in some case it is useful to store them precomputed?

from mace.

I now noticed that this change breaks compatibility with loading models.

E.g.: if you load a model that stored the U matrices in a version of mace that has them set to persistent=False this will result in an error. And the same happens in the opposite case.

from mace.

I now noticed that this change breaks compatibility with loading models.

E.g.: if you load a model that stored the U matrices in a version of mace that has them set to persistent=False this will result in an error. And the same happens in the opposite case.

This could maybe be fixed with non-strict loading. The torch load function has a keyword strict=False, but it might be better to just break the backwards compatibility or explicitly remove the U matrices from old checkpoints.

from mace.

@ilyes319 do you think you can make them persistent=False? To load an old model on the new implementation it would just be a matter of "cleaning" the checkpoint file, i.e. removing the matrices from it.

from mace.

I wonder how this interacts with torchscript though and libtorch. I guess the safest would be to make an argument and keep the default to true. Would this alright?

from mace.

Yes, if it can be configured from an argument of the SymmetricContraction module (not just Contraction) I think it would be fine for us 👍

from mace.

Could we add this? :) (a persistent_U_matrices or something similar argument to SymmetricContraction that defaults to True)

I can submit a PR.

from mace.