如何从TensorFlow Federated中提取聚合梯度？

在Weighted FedAvg算法中，服务端并非直接聚合客户端梯度，而是聚合客户端本地训练后的权重与初始全局权重的差值（即本地更新），再将加权平均后的更新应用到全局权重上。你可以通过以下两种方式获取聚合后的梯度/更新：

方法1：计算聚合前后的权重差值（推荐） #

这种方法可以得到服务端实际应用到全局权重上的最终更新，步骤如下：

# 初始化联邦训练状态
state = iterative_process.initialize()

# 记录聚合前的全局权重
pre_global_weights = iterative_process.get_model_weights(state)

# 运行一轮联邦训练（client_data为你的客户端数据集）
state, round_metrics = iterative_process.next(state, client_data)

# 记录聚合后的全局权重
post_global_weights = iterative_process.get_model_weights(state)

# 递归计算权重差值，得到聚合后的更新（等效于聚合梯度的作用）
aggregated_updates = tf.nest.map_structure(
    lambda pre_w, post_w: post_w - pre_w,
    pre_global_weights,
    post_global_weights
)

# 示例：查看第一层全连接层的权重更新
print("第一层Dense权重更新：", aggregated_updates.trainable[0])
# 查看第一层全连接层的偏置更新
print("第一层Dense偏置更新：", aggregated_updates.trainable[1])

说明 #

• tf.nest.map_structure用于递归处理嵌套的权重结构（模型权重由多个层的权重/偏置组成，是嵌套Tensor结构）
• 如果服务端使用了优化器（比如你代码中的Adam），这个aggregated_updates是经过优化器调整后的最终权重变化，而非原始的客户端更新加权平均。

方法2：获取原始客户端更新的加权平均（需自定义逻辑） #

如果需要获取未经过服务端优化器处理的原始聚合客户端更新（即所有客户端本地delta的加权平均），需要自定义FedAvg的服务端更新逻辑，或者拆解build_weighted_fed_avg的内部实现：

def server_update_fn(server_state, aggregated_client_deltas):
    # 这里可以直接获取aggregated_client_deltas，即原始聚合更新
    print("原始聚合客户端更新：", aggregated_client_deltas)
    # 继续执行原有的服务端优化逻辑
    updated_weights = tf.nest.map_structure(
        lambda w, delta: server_state.optimizer.apply_gradients(zip([delta], [w])),
        server_state.model_weights,
        aggregated_client_deltas
    )
    return server_state._replace(model_weights=updated_weights)

# 基于自定义server_update_fn构建迭代过程
# （注：此代码仅为示例，需要结合tff.learning.algorithms的低级组件完整实现）

不过这种方式需要对TFF的联邦学习流程有较深入的理解，大多数场景下方法1已经满足需求。

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