如何加载ServerState.optimizer_state以继续在Tensorflow Federated中进行培训



TFF是否有任何类似于模型权重的方法来保存和加载优化器状态。对于模型权重,有ModelWeights.assign_weights_to()tff.learning.state_with_new_model_weights()函数,是否有保存和加载优化器状态的方法,尤其是在使用SGD以外的服务器端优化器时。

我找不到任何东西来保存和加载优化器的状态。

对Keith的回答进行回顾:检查点管理器在谷歌的联合研究库中不再可用。它已经被上传到TFF(请在GitHub上查看(。

您可以通过tensorflow-federated-nightlypip包访问它,也可以通过克隆存储库来访问它。

代码本质上是委托给tf.saved_model.save的,所以您也可以简单地使用它。

这应该可以通过TFF的tff.simulation.FileCheckpointManager实现。特别是,Google Research的federatedrepo中的用法最初是为了支持在服务器上使用学习率调度和自适应优化时从检查点重新启动而编写的,对于该应用程序,正确恢复优化器状态至关重要。

只要tff.templates.IterativeProcess返回适当的优化器状态,简单地使用开箱即用的FileCheckpointManager就可以了。