☰
用Python实现本地模拟横向联邦学习:一台笔记本就能跑通的隐私计算入门
2026/10/14 12:00:17 网站建设 项目流程

简介:这是一份面向机器学习与分布式系统学习者的实战资料,围绕用Python在本地模拟横向联邦学习展开,适合希望理解联邦学习原理、动手搭建多客户端协作训练环境的开发者与高校学生。资源包共25个文件,以py脚本、pyc缓存、xml配置及CIFAR-10数据批次文件为主,涵盖服务端与客户端逻辑、模型定义、数据集加载及项目配置,压缩包约302.68MB。已有723人学习下载。通过server.py、client.py、models.py、datasets.py等模块,读者可完整走通本地训练、参数上传、服务器聚合与全局模型广播的流程,并借助conf.json与main.py快速调整实验参数;配套的CIFAR-10数据可直接用于复现实验,理解横向联邦学习在隐私保护前提下的协作训练机制,同时为后续引入gRPC通信、差分隐私等扩展打下工程基础。

1. 用Python实现本地模拟横向联邦学习:一台笔记本就能跑通的隐私计算入门

你可能已经在各种技术社区刷到过「联邦学习」这个词,但真要动手时发现,大部分教程一上来就让你部署分布式集群、配 gRPC 通信、搭参数服务器,光环境就能劝退一半人。其实横向联邦学习的核心逻辑并不复杂:多个客户端各自持有本地数据,在不上传原始数据的前提下,只交换模型参数或梯度,由中心节点聚合后下发更新。这件事完全可以用 Python 在一台笔记本上模拟出来,不需要任何真实网络通信,也不需要多台机器。这篇文章面向的是想快速理解联邦学习到底怎么运转、想先跑通一个最小闭环再决定要不要深入的人。我会用 NumPy 和 PyTorch 从零搭一个本地模拟框架,把数据划分、客户端训练、参数聚合、多轮迭代全部串起来,每一步都有可复制的代码和参数说明。跑完这一套,你对 FedAvg 的理解不会再停留在论文摘要层面。

2. 横向联邦学习的数据划分与客户端构造:为什么 IID 和非 IID 是两回事

2.1 横向联邦的适用场景与数据切分逻辑

横向联邦学习(Horizontal Federated Learning)的前提是:所有参与方共享同一套特征空间,但各自持有的样本不同。比如三家不同的机构都在做用户行为预测,特征字段一样,但用户群体不重叠。这种情况下,数据是按样本维度横向切开的,所以叫「横向」。

在本地模拟时,最直接的做法就是把一份完整数据集按样本切给 N 个客户端。但这里有个关键选择:切分方式决定了你模拟出来的难度。如果随机均匀切分,每个客户端的数据分布和全局分布基本一致,这叫 IID(独立同分布)划分。如果按标签类别切分,比如客户端 A 只有类别 0 和 1 的数据,客户端 B 只有类别 2 和 3 的数据,这就是非 IID 划分,也是真实场景中更常见的情况。

我一般会两种都跑一遍,因为 IID 下模型收敛快、准确率高,容易让人产生「联邦学习很简单」的错觉;非 IID 下才会暴露梯度漂移、客户端模型差异大、聚合后全局模型震荡这些真正的问题。

import numpy as np import torch from torch.utils.data import DataLoader, TensorDataset def split_iid(X, y, num_clients, seed=42): """IID 划分:随机均匀分配给每个客户端""" rng = np.random.default_rng(seed) indices = rng.permutation(len(X)) splits = np.array_split(indices, num_clients) clients = [] for idx in splits: clients.append((X[idx], y[idx])) return clients def split_non_iid(X, y, num_clients, classes_per_client=2, seed=42): """非 IID 划分:每个客户端只拿部分类别的数据""" rng = np.random.default_rng(seed) num_classes = len(np.unique(y)) class_indices = {c: np.where(y == c)[0] for c in range(num_classes)} for c in class_indices: rng.shuffle(class_indices[c]) client_classes = [] all_classes = list(range(num_classes)) rng.shuffle(all_classes) for i in range(num_clients): start = (i * classes_per_client) % num_classes assigned = [all_classes[(start + j) % num_classes] for j in range(classes_per_client)] client_classes.append(assigned) clients = [] for assigned in client_classes: idx_list = [] for c in assigned: if len(class_indices[c]) > 0: take = max(1, len(class_indices[c]) // (num_clients // classes_per_client + 1)) idx_list.extend(class_indices[c][:take]) class_indices[c] = class_indices[c][take:] idx = np.array(idx_list) if len(idx) == 0: idx = rng.choice(len(X), size=32, replace=False) clients.append((X[idx], y[idx])) return clients

split_iid的逻辑很直白:打乱全局索引后均分。split_non_iid稍微绕一点,它先把每个类别的样本索引单独拿出来打乱,然后给每个客户端分配classes_per_client个类别,从对应类别的池子里取一部分。classes_per_client这个参数控制非 IID 程度——设成 2 意味着每个客户端只见过 2 个类,设成总类别数就退化成 IID。seed保证可复现,调参对比时不会因为随机性导致结论不可信。

2.2 客户端本地训练的封装

每个客户端本质上就是一个持有本地数据、能独立跑若干轮 SGD 的模块。我习惯把它写成一个类,方便后续扩展(比如加差分隐私、加客户端选择策略)。

class FederatedClient: def __init__(self, client_id, X, y, batch_size=32, lr=0.01): self.client_id = client_id self.X = torch.tensor(X, dtype=torch.float32) self.y = torch.tensor(y, dtype=torch.long) self.dataset = TensorDataset(self.X, self.y) self.loader = DataLoader(self.dataset, batch_size=batch_size, shuffle=True) self.lr = lr self.model = None def set_model(self, model): """接收全局模型参数""" self.model = model def local_train(self, epochs=1): """本地训练,返回更新后的参数和样本数""" self.model.train() optimizer = torch.optim.SGD(self.model.parameters(), lr=self.lr) criterion = torch.nn.CrossEntropyLoss() for _ in range(epochs): for bx, by in self.loader: optimizer.zero_grad() output = self.model(bx) loss = criterion(output, by) loss.backward() optimizer.step() return self.model.state_dict(), len(self.X)

local_train返回两个东西:模型参数字典和本地样本数。样本数在后面聚合时当权重用——样本多的客户端对全局模型的贡献更大,这是 FedAvg 的标准做法。epochs控制本地训练轮数,设太大容易导致客户端模型过度拟合本地数据,聚合后反而变差,这个坑后面会细说。

3. FedAvg 聚合的 Python 实现:从参数平均到加权策略

3.1 聚合函数的写法与权重设计

FedAvg 的核心就一句话:按样本数加权平均各客户端的模型参数。听起来简单,但实现时有几个细节容易翻车——比如哪些参数需要聚合、哪些不需要;比如 BatchNorm 的 running_mean 和 running_var 要不要一起平均。

def federated_averaging(client_states, client_sizes): """FedAvg:按样本数加权平均模型参数""" total_samples = sum(client_sizes) avg_state = {} for key in client_states[0].keys(): avg_state[key] = sum( client_states[i][key] * (client_sizes[i] / total_samples) for i in range(len(client_states)) ) return avg_state

这段代码假设所有客户端的模型结构完全一致,state_dict的 key 一一对应。client_sizes[i] / total_samples就是第 i 个客户端的聚合权重。如果所有客户端样本数相同,就退化成等权平均。

有一个容易忽略的点:PyTorch 的state_dict里包含num_batches_tracked这类整型 buffer,直接做浮点乘法会报类型错误。稳妥做法是在聚合前过滤掉非浮点参数,或者对整型参数单独处理(通常取第一个客户端的值就行)。

def federated_averaging_safe(client_states, client_sizes): """安全版:跳过非浮点类型的 buffer""" total_samples = sum(client_sizes) avg_state = {} for key in client_states[0].keys(): if not torch.is_floating_point(client_states[0][key]): avg_state[key] = client_states[0][key].clone() continue avg_state[key] = sum( client_states[i][key].float() * (client_sizes[i] / total_samples) for i in range(len(client_states)) ) return avg_state

3.2 完整的联邦训练循环

把客户端、聚合、全局模型串起来,就是一个最小的联邦学习系统。下面这段代码可以直接跑,用的是 sklearn 的 digits 数据集(手写数字,10 类,1797 个样本),足够小,笔记本 CPU 几秒就能跑完一轮。

import torch.nn as nn from sklearn.datasets import load_digits from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler # 数据准备 digits = load_digits() X, y = digits.data, digits.target scaler = StandardScaler() X = scaler.fit_transform(X) X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # 定义模型 class SimpleNet(nn.Module): def __init__(self, input_dim=64, num_classes=10): super().__init__() self.fc = nn.Sequential( nn.Linear(input_dim, 128), nn.ReLU(), nn.Linear(128, num_classes) ) def forward(self, x): return self.fc(x) # 联邦训练主循环 def run_federated(X_train, y_train, X_test, y_test, num_clients=5, num_rounds=20, local_epochs=2, non_iid=False): if non_iid: client_data = split_non_iid(X_train, y_train, num_clients) else: client_data = split_iid(X_train, y_train, num_clients) global_model = SimpleNet() clients = [] for i, (cx, cy) in enumerate(client_data): clients.append(FederatedClient(i, cx, cy, batch_size=32, lr=0.01)) test_tensor = torch.tensor(X_test, dtype=torch.float32) test_label = torch.tensor(y_test, dtype=torch.long) for rnd in range(num_rounds): states, sizes = [], [] for client in clients: # 每个客户端从全局模型出发 import copy local_model = copy.deepcopy(global_model) client.set_model(local_model) state, size = client.local_train(epochs=local_epochs) states.append(state) sizes.append(size) # 聚合 new_state = federated_averaging_safe(states, sizes) global_model.load_state_dict(new_state) # 评估 global_model.eval() with torch.no_grad(): pred = global_model(test_tensor).argmax(dim=1) acc = (pred == test_label).float().mean().item() print(f"Round {rnd+1:02d} | Test Acc: {acc:.4f}") return global_model # 跑 IID print("=== IID 划分 ===") run_federated(X_train, y_train, X_test, y_test, non_iid=False) # 跑非 IID print("=== 非 IID 划分 ===") run_federated(X_train, y_train, X_test, y_test, non_iid=True)

num_rounds是全局通信轮数,local_epochs是每轮每个客户端本地训练的 epoch 数。copy.deepcopy(global_model)这一步很关键——每个客户端必须从同一份全局参数出发,不能共享同一个模型对象,否则客户端 A 的训练会污染客户端 B 的起点。

跑 IID 划分时,通常 10 到 15 轮就能到 90% 以上的准确率。非 IID 划分下,同样的轮数可能只有 70% 到 80%,而且曲线会抖得厉害。这不是代码写错了,是非 IID 本身的特性。

4. 本地模拟联邦学习的避坑指南:五个血泪教训

4.1 客户端模型没有深拷贝导致结果全错

现象:训练几轮后准确率不升反降,或者所有客户端返回的 state_dict 完全一样。

原因:多个客户端共享了同一个模型对象。Python 里对象赋值是引用传递,client.set_model(global_model)如果直接传引用,所有客户端操作的是同一个模型,本地训练变成串行叠加。

解决:每个客户端训练前用copy.deepcopy复制一份全局模型,训练完只取 state_dict,不保留模型对象。

4.2 本地 epoch 设太大导致全局模型震荡

现象:非 IID 场景下,local_epochs从 1 调到 5 后,全局准确率反而下降,loss 曲线剧烈波动。

原因:本地训练轮数越多,客户端模型越偏向本地数据分布。非 IID 下各客户端的最优方向差异很大,聚合时相互抵消。

解决:非 IID 场景把local_epochs控制在 1 到 2,或者引入客户端选择策略(每轮只随机选一部分客户端参与)。FedProx 那类加正则项的方法也是为这个场景设计的。

4.3 忘记过滤非浮点 buffer 导致聚合报错

现象:federated_averaging执行时报RuntimeError: result type Long can't be cast to the desired output type Float。

原因:state_dict里混有num_batches_tracked这类整型张量,浮点加权平均对它们不适用。

解决:聚合前判断torch.is_floating_point,非浮点参数直接取第一个客户端的值克隆。

4.4 数据没有归一化导致训练不收敛

现象:IID 划分下跑了 50 轮准确率还在 20% 左右,loss 几乎不变。

原因:digits 数据集原始像素值范围是 0 到 16,直接喂给网络会导致梯度尺度不一致,SGD 很难找到下降方向。

解决:训练前用StandardScaler做零均值单位方差归一化。这个坑和联邦学习本身无关,但因为是本地模拟,很多人会忽略数据预处理。

4.5 测试集泄露到客户端本地数据

现象:全局模型准确率高得离谱,但换一份测试集就崩了。

原因:划分客户端数据时不小心把测试集样本混进去了,或者用全量数据做了 scaler 拟合。

解决:先切分训练集和测试集,scaler 只在训练集上 fit,客户端划分只在训练集内部进行。测试集从头到尾不参与任何训练环节。

5. 进阶技巧:用非 IID 程度和通信轮数做消融实验

跑通最小闭环之后,最有价值的动作是做消融实验——固定其他变量,只调一个参数,看全局模型的表现怎么变。我一般会重点盯两个维度:非 IID 程度和通信轮数。

非 IID 程度可以用classes_per_client来控制。设成 10(等于总类别数)就是 IID,设成 2 就是极端非 IID。中间可以取 3、5、8,画一条准确率随非 IID 程度变化的曲线。你会发现一个规律:当classes_per_client降到 3 以下时,准确率下降速度明显加快,而且轮间方差变大。这个拐点对判断你的场景能不能用联邦学习很关键。

# 消融实验:不同非 IID 程度下的最终准确率 results = {} for cpc in [2, 3, 5, 8, 10]: print(f"\n--- classes_per_client = {cpc} ---") model = run_federated(X_train, y_train, X_test, y_test, num_clients=5, num_rounds=15, local_epochs=2, non_iid=(cpc < 10)) # 记录最后一轮准确率 results[cpc] = acc # acc 从 run_federated 内部打印获取,实际使用时可以改成返回值

通信轮数的消融更直接:把num_rounds从 5 跑到 50,每隔 5 轮记录一次准确率。IID 下通常 15 到 20 轮就收敛了,非 IID 下可能需要 40 轮以上,而且最终准确率上限更低。这个实验能帮你估算真实部署时需要多少通信开销。

还有一个值得试的技巧是客户端采样。每轮不选全部客户端,而是随机选一部分(比如 50%),这更接近真实场景——有些客户端可能掉线或忙碌。实现上就是在训练循环里加一行随机采样:

import random def select_clients(clients, fraction=0.5): num_selected = max(1, int(len(clients) * fraction)) return random.sample(clients, num_selected) # 在主循环里替换: # for client in clients: for client in select_clients(clients, fraction=0.5): ...

采样比例越低,每轮通信开销越小,但收敛需要的轮数越多。这个权衡在真实部署里直接对应带宽成本和训练时间的取舍。

我自己踩过最深的一个坑是:一开始觉得非 IID 下准确率低是代码有 bug,花了整整一个下午逐行排查,最后发现把local_epochs从 2 改成 1 就正常了。联邦学习里很多「玄学」现象,根因都在本地训练和全局聚合之间的拉扯。建议你跑通之后,先把local_epochs、lr、num_clients这三个参数各扫一遍,观察准确率曲线的变化,比看十篇论文都管用。希望帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询