联邦学习（二）：隐私保护的分布式机器学习

# FedAvg 核心算法实现
import numpy as np
from typing import List, Tuple

def fedavg_round(
    global_weights: List[np.ndarray],
    client_weights: List[List[np.ndarray]],
    client_data_sizes: List[int],
) -> List[np.ndarray]:
    """执行一轮 FedAvg 聚合。
    
    Args:
        global_weights: 当前全局模型权重列表（每层一个数组）
        client_weights: 各客户端更新后的权重，格式为 [client_i][layer_j]
        client_data_sizes: 各客户端的数据量（用于加权）
    """
    total_size = sum(client_data_sizes)
    new_global = []
    
    for layer_idx in range(len(global_weights)):
        # 按数据量加权平均
        weighted_sum = np.zeros_like(global_weights[layer_idx])
        for client_idx, cw in enumerate(client_weights):
            weight = client_data_sizes[client_idx] / total_size
            weighted_sum += weight * cw[layer_idx]
        new_global.append(weighted_sum)
    
    return new_global


def client_update(
    local_weights: List[np.ndarray],
    local_data,
    local_labels,
    epochs: int = 5,
    learning_rate: float = 0.01,
) -> List[np.ndarray]:
    """客户端本地训练。
    
    每个客户端在本地数据上执行多轮 SGD，
    这是 FedAvg 减少通信轮数的关键设计。
    """
    # 简化的 SGD 实现（实际使用 PyTorch/JAX）
    for epoch in range(epochs):
        # 计算梯度
        gradients = compute_gradients(local_weights, local_data, local_labels)
        # 更新权重
        for layer_idx in range(len(local_weights)):
            local_weights[layer_idx] -= learning_rate * gradients[layer_idx]
    
    return local_weights


# 模拟一轮联邦学习
n_clients = 10
n_layers = 4
global_weights = [np.random.randn(100) for _ in range(n_layers)]

# 各客户端本地训练
client_updates = []
data_sizes = []
for i in range(n_clients):
    local_w = [w.copy() for w in global_weights]
    local_w = client_update(local_w, f"data_{i}", f"labels_{i}")
    client_updates.append(local_w)
    data_sizes.append(np.random.randint(500, 5000))

# 服务器聚合
global_weights = fedavg_round(global_weights, client_updates, data_sizes)
print(f"完成一轮 FedAvg 聚合，{n_clients} 个客户端参与")

# DP-FedAvg：差分隐私联邦学习
import numpy as np

class DPClient:
    """支持差分隐私的联邦学习客户端"""
    
    def __init__(
        self, 
        client_id: int,
        noise_multiplier: float = 1.0,
        max_grad_norm: float = 1.0,
    ):
        self.client_id = client_id
        self.noise_multiplier = noise_multiplier  # 噪声倍数 σ
        self.max_grad_norm = max_grad_norm        # 梯度裁剪阈值 C
    
    def compute_private_gradient(
        self, 
        model, 
        data, 
        labels,
        batch_size: int = 32,
    ) -> np.ndarray:
        """计算带差分隐私保护的梯度。
        
        步骤：
        1. 计算每个样本的梯度
        2. 按 max_grad_norm 裁剪每个样本梯度
        3. 聚合裁剪后的梯度
        4. 添加高斯噪声
        """
        # 计算每个样本梯度（实际框架中自动完成）
        per_sample_grads = compute_per_sample_gradients(
            model, data, labels
        )
        
        # 梯度裁剪
        norms = np.linalg.norm(per_sample_grads, axis=1, keepdims=True)
        clipped_grads = per_sample_grads * np.minimum(
            1.0, self.max_grad_norm / (norms + 1e-8)
        )
        
        # 聚合
        avg_grad = clipped_grads.mean(axis=0)
        
        # 添加高斯噪声
        noise = np.random.normal(
            0, 
            self.noise_multiplier * self.max_grad_norm, 
            avg_grad.shape
        )
        
        return avg_grad + noise
    
    def get_privacy_budget(
        self, n_rounds: int, n_clients: int, delta: float = 1e-5
    ) -> float:
        """使用高级组合定理估算 ε 隐私预算。
        
        简化版 RDP → (ε, δ)-DP 转换
        """
        q = 1.0 / n_clients  # 采样率
        sigma = self.noise_multiplier
        # 简化的隐私预算估算（实际使用 opacus 库）
        epsilon = np.sqrt(
            2 * n_rounds * np.log(1/delta) * (q2) / (sigma2)
        )
        return epsilon


# 使用示例
client = DPClient(client_id=0, noise_multiplier=0.5, max_grad_norm=1.0)
# private_grad = client.compute_private_gradient(model, data, labels)
budget = client.get_privacy_budget(n_rounds=100, n_clients=50)
print(f"100 轮训练后的隐私预算 ε ≈ {budget:.2f}")

# Flower 框架示例：构建一个联邦学习系统
import flwr as fl
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# 1. 定义模型
class SimpleNet(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, num_classes: int):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, num_classes)
    
    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

# 2. 定义 Flower Client
class FlowerClient(fl.client.NumPyClient):
    def __init__(self, model, train_loader, val_loader):
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
    
    def get_parameters(self, config):
        return [v.cpu().numpy() for v in self.model.state_dict().values()]
    
    def set_parameters(self, parameters):
        state_dict = {
            k: torch.tensor(v) 
            for k, v in zip(self.model.state_dict().keys(), parameters)
        }
        self.model.load_state_dict(state_dict)
    
    def fit(self, parameters, config):
        """本地训练"""
        self.set_parameters(parameters)
        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.01)
        criterion = nn.CrossEntropyLoss()
        
        self.model.train()
        for epoch in range(3):  # 本地 3 个 epoch
            for batch_x, batch_y in self.train_loader:
                batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
                optimizer.zero_grad()
                loss = criterion(self.model(batch_x), batch_y)
                loss.backward()
                optimizer.step()
        
        return self.get_parameters(config={}), len(self.train_loader.dataset), {}
    
    def evaluate(self, parameters, config):
        """本地评估"""
        self.set_parameters(parameters)
        self.model.eval()
        
        correct, total, loss = 0, 0, 0.0
        criterion = nn.CrossEntropyLoss()
        with torch.no_grad():
            for batch_x, batch_y in self.val_loader:
                batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
                outputs = self.model(batch_x)
                loss += criterion(outputs, batch_y).item()
                correct += (outputs.argmax(1) == batch_y).sum().item()
                total += batch_y.size(0)
        
        accuracy = correct / total
        return float(loss), total, {"accuracy": accuracy}

# 3. 启动 Flower 服务器
# strategy = fl.server.strategy.FedProx(
#     fraction_fit=0.1,      # 每轮选择 10% 客户端
#     fraction_evaluate=0.1,
#     proximal_mu=0.1,       # FedProx 的 μ 参数
# )
# fl.server.start_server(
#     server_address="0.0.0.0:8080",
#     config=fl.server.ServerConfig(num_rounds=50),
#     strategy=strategy,
# )