联邦学习：是否有工具库可获取每轮各本地模型权重以实现客户端聚类？

自定义训练流程（最直接方案） #

如果是自己搭建的联邦训练逻辑，无需依赖框架封装，直接在客户端训练函数末尾返回权重即可：

• 客户端训练代码示例（PyTorch）：

def local_train(model, train_loader, optimizer, loss_fn, local_epochs):
    model.train()
    for _ in range(local_epochs):
        for x, y in train_loader:
            optimizer.zero_grad()
            pred = model(x)
            loss = loss_fn(pred, y)
            loss.backward()
            optimizer.step()
    # 训练完成后返回模型权重
    return model.state_dict()

• 服务器端收集所有客户端权重：

client_weights = []
for client in client_list:
    weights = local_train(client.model, client.train_data, client.optimizer, client.loss_fn, local_epochs)
    client_weights.append(weights)
# 后续可将权重展平为一维向量，用K-Means/DBSCAN等算法聚类

TensorFlow Federated (TFF) 场景 #

TFF中可通过自定义客户端计算逻辑，让客户端返回训练后的权重：

import tensorflow_federated as tff
import tensorflow as tf

# 定义基础Keras模型
def create_keras_model():
    return tf.keras.Sequential([
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

# 封装客户端训练+权重返回逻辑
@tff.tf_computation(tff.SequenceType(tf.float32), tff.ModelWeightsType(create_keras_model()))
def client_train_fn(data, initial_weights):
    model = tff.learning.from_keras_model(
        keras_model=create_keras_model(),
        input_spec=data.element_spec,
        loss=tf.keras.losses.SparseCategoricalCrossentropy()
    )
    model_weights = tff.learning.ModelWeights.from_model(model)
    model_weights.assign(initial_weights)
    
    # 本地训练步骤
    for batch in data:
        x, y = batch
        with tf.GradientTape() as tape:
            pred = model(x)
            loss = model.loss(y, pred)
        grads = tape.gradient(loss, model.trainable_variables)
        model.optimizer.apply_gradients(zip(grads, model.trainable_variables))
    
    # 返回训练后的本地权重
    return model_weights

# 联邦训练循环中收集权重
for round_num in range(total_rounds):
    global_weights = training_process.get_model_weights(train_state)
    # 分发全局权重到客户端，执行训练并收集本地权重
    client_weights = tff.federated_eval(client_train_fn, tff.federated_broadcast(global_weights), client_dataset)
    # 基于client_weights执行聚类操作

FedML 框架场景 #

FedML可通过自定义客户端Trainer导出本地权重：

from fedml.core.client.client_trainer import ClientTrainer

class WeightReturningTrainer(ClientTrainer):
    def train(self, train_data, device):
        # 执行默认训练流程
        super().train(train_data, device)
        # 返回训练后的模型权重
        return self.model.state_dict()

# 初始化客户端管理器时使用自定义Trainer
client_manager = FedMLClientManager(client_list, WeightReturningTrainer)

# 每轮训练后收集权重
for round in range(total_rounds):
    training_process.train()
    # 遍历客户端获取所有本地权重
    client_weights = [client.trainer.train(client.train_data, client.device) for client in client_manager.clients]

通用思路 #

无论使用哪种框架，核心逻辑都是在客户端本地训练完成的节点，将权重传递回服务器端：

• 对于封装度高的框架，优先查找官方文档中「客户端训练钩子」「自定义计算返回值」类的API；
• 收集到权重后，需将各层参数展平为一维向量（比如PyTorch中把state_dict的values转为numpy后拼接），再用聚类算法完成客户端分组。

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