如何保存与重载TensorFlow Federated模型并更新迭代状态？

核心思路 #

TFF的ServerState是结构化的TensorFlow对象，包含模型权重、聚合状态等核心数据，需借助TensorFlow的检查点机制（tf.train.Checkpoint）进行序列化与恢复，确保状态结构完全匹配即可重载复用。

保存方法 #

方案一：完整保存整个ServerState #

适合需要保留所有迭代状态（包括聚合统计、广播状态等）的场景：

import tensorflow as tf

# 绑定状态到检查点对象
checkpoint = tf.train.Checkpoint(server_state=state)
# 保存到指定目录，可多次调用实现增量保存
save_path = checkpoint.save('./tff_server_checkpoints')

方案二：仅保存模型可训练参数 #

若仅关注模型权重，可单独提取保存，减少存储体积：

# 提取模型可训练参数列表
trainable_weights = state.model.trainable
# 绑定参数到检查点
checkpoint = tf.train.Checkpoint(model_weights=trainable_weights)
checkpoint.save('./tff_model_weights')

重载方法 #

对应方案一：恢复完整ServerState #

需先初始化与原迭代过程结构完全一致的新状态，再通过检查点恢复：

# 初始化新的迭代过程（需与原定义完全相同）
new_iterative_process = ... # 例如：tff.learning.build_federated_averaging_process(...)
new_state = new_iterative_process.initialize()

# 绑定新状态到检查点
checkpoint = tf.train.Checkpoint(server_state=new_state)
# 恢复最新检查点
latest_checkpoint = tf.train.latest_checkpoint('./tff_server_checkpoints')
# 断言所有变量恢复完成，避免结构不匹配
checkpoint.restore(latest_checkpoint).assert_consumed()

对应方案二：恢复模型可训练参数 #

初始化新状态后，直接将恢复的权重赋值给新状态的模型参数：

new_state = new_iterative_process.initialize()

# 绑定新状态的模型参数到检查点
checkpoint = tf.train.Checkpoint(model_weights=new_state.model.trainable)
latest_checkpoint = tf.train.latest_checkpoint('./tff_model_weights')
checkpoint.restore(latest_checkpoint).assert_consumed()

关键注意事项 #

• 结构一致性：新迭代过程的模型结构、聚合策略、状态组件必须与原定义完全一致，否则检查点恢复会失败。
• 版本兼容：若需跨TFF版本复用，优先选择保存模型可训练参数，ServerState的内部结构可能随TFF版本更新变化。
• 验证恢复结果：可通过TensorFlow断言工具验证参数是否正确恢复：

tf.debugging.assert_equal(state.model.trainable[0], new_state.model.trainable[0])

内容的提问来源于stack exchange，提问作者Los