联邦学习场景下对PyTorch参数做联邦平均后如何更新全局模型参数？

1 核心问题定位 #

当前代码存在3个核心问题导致报错与逻辑错误：

• 参数收集函数RPC_get_parameters逻辑错误：遍历模型参数时直接return，仅能返回模型第一层的参数，无法获取全部层的参数
• 参数平均逻辑错误：未按层对应累加不同节点的参数，同时计数逻辑global_count += len(global_sum)完全不符合联邦平均的计算规则
• 参数更新逻辑错误：优化器需要绑定模型本身的可学习参数（叶张量），直接传入计算得到的平均张量会触发非叶张量报错

2 修正参数收集逻辑 #

首先修改RPC_get_parameters函数，返回所有层的参数，且统一转成无梯度的CPU张量避免RPC传输问题：

def RPC_get_parameters(data, model):
    """
    Get all parameters from nodes
    """
    param_list = []
    with torch.no_grad():
        for param in model.parameters():
            # detach去除梯度，转CPU避免设备不匹配
            param_list.append(param.detach().cpu())
    return {"parameters": param_list}

3 修正联邦平均逻辑 #

按层对应累加所有节点的参数，再按节点数（如果是加权联邦平均可替换为样本权重）计算平均：

results = client.get_results(task_id=task.get("id"))
num_nodes = len(results)

# 初始化每层的参数和，和第一层的参数形状对齐
global_param_sum = [torch.zeros_like(p) for p in results[0]["parameters"]]

for output in results:
    node_params = output["parameters"]
    # 按层累加参数
    for layer_idx in range(len(global_param_sum)):
        global_param_sum[layer_idx] += node_params[layer_idx]

# 计算平均参数
averaged_params = [layer_sum / num_nodes for layer_sum in global_param_sum]

4 正确更新全局模型参数 #

不要直接把平均参数传给优化器，而是先把平均参数赋值给全局模型的对应参数的data属性，优化器仍然绑定全局模型的参数即可：

# 假设已提前初始化全局模型 global_model
with torch.no_grad():
    for model_param, avg_param in zip(global_model.parameters(), averaged_params):
        # 直接赋值给参数的data属性，不会修改计算图，保持参数为叶张量
        model_param.data.copy_(avg_param.to(model_param.device))

# 优化器初始化绑定全局模型的参数即可，无需修改传入对象
optimizerD = optim.SGD(global_model.parameters(), lr=0.01, momentum=0.5)

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