How to save Tensorflow Federated Model

Question

How can I easily save a Tensorflow Federated model? (state)

A few months ago I was using this solution after importing ServerState and FileCheckPointManager and it worked:

# Create the checkpoint Manager
ckpt_manager = FileCheckpointManager(root_dir=checkpoint_dir)
# Save checkpoint for round N
ckpt_manager.save_checkpoint(ServerState.from_tff_result(state), round_num=NUM_ROUNDS)

But now this solution no longer works because ServerState does not contain from_tff_result methotd anymore.

AttributeError: type object 'ServerState' has no attribute 'from_tff_result'

Also using the old version of ServerState where the metohd was included i get:

TypeError: Expected tensorflow_federated.python.common_libs.structure.Struct, found tensorflow_federated.python.learning.model_utils.ModelWeights.

How can I easly save my federated model?

Answer 1

After a recent change to prevent Struct (formerly AnonymousTuple ) from being returned to users, there should no longer be any need for a method like ServerState.from_tff_result(state) ; state should already be a ServerState .

TLDR: if you update your example to simply read:

# Create the checkpoint Manager
ckpt_manager = FileCheckpointManager(root_dir=checkpoint_dir)
# Save checkpoint for round N
ckpt_manager.save_checkpoint(state, round_num=NUM_ROUNDS)

it should just work.