Saving models to .h5 in tf>=2.8 fails

[andreped]

Saving models to .h5 in tf>=2.8 fails

[andreped]

Hmm, is it because of the name attribute not being set for this tf.Variable maybe:
https://github.com/andreped/GradientAccumulator/blob/main/gradient_accumulator/GAModelWrapper.py#L11

Could you try setting a unique name for this variable and running one epoch, to see if the model checkpoint is able to save (if that is where it fails)?

[andreped]

I was able to reproduce the error, by using .h5 when saving the model in a unit test:
https://github.com/andreped/GradientAccumulator/runs/7013532197?check_suite_focus=true

[andreped]

Just did a sanity check (see run here) where I removed the GA model wrapping and saved with .h5, and then it works, so the issue definitely comes from the GAModelWrapper somewhere.

[andreped]

Hmm, seems like this works just fine in TF 2.6.2, as seen from this CI job:
https://github.com/andreped/GradientAccumulator/runs/7013531511?check_suite_focus=true

But for newer versions of TF it fails. I read that they encourage us to use the SavedModel format instead of the old possibly-soon-deprecated HDF5 format.

So perhaps we should just avoid having HDF5 format support here, and instead focus on getting the other stuff working properly? I will refactor the stuff I did now and remove the .h5 test.

[andreped]

Tried lots of more stuff. Model wrapping just breaks with the old HDF5 format. That makes sense as this new train step overloading was introduced in TF 2.2, where SavedModel was the new default model format.

There are some pros/cons between SavedModel/HDF5, which you could read about here:
https://keras.io/guides/serialization_and_saving/#format-limitations

But most importantly, I believe when stored in HDF5 format, the model wrapping just don't work when loading. I believe this works with the SavedModel format because the model wrapper is not saved, as mentioned in the link above. So that format does not have the same issue. I tried having custom objects for this, but this is a lot more tricky with model wrappers, compared to regular layers.

In the end, just use SavedModel. HDF5 is deprecated. I will add a comment about this in the README. Closing for now.