┌─────────────────┐
         │    中央服务器    │
         │  (Aggregator)   │
         └────────┬────────┘
                  │ 模型分发 / 更新聚合
         ┌────────┼────────┐
         ↓        ↓        ↓
    ┌─────────┐ ┌─────────┐ ┌─────────┐
    │客户端 1  │ │客户端 2  │ │客户端 3  │
    │(用本地  │ │(用本地  │ │(用本地  │
    │数据训练)│ │数据训练)│ │数据训练)│
    └─────────┘ └─────────┘ └─────────┘

客户端 1: [特征1, 特征2, 特征3] x [样本 1~1000]
客户端 2: [特征1, 特征2, 特征3] x [样本 1001~2000]
客户端 3: [特征1, 特征2, 特征3] x [样本 2001~3000]

客户端 A(银行): [金融特征] x [用户 1~10000]
客户端 B(医院): [医疗特征] x [用户 1~10000]

服务器执行:
  初始化 w_0
  for 轮次 t = 1, 2, ..., T:
    m = max(C x K, 1)  // C: 参与比例, K: 客户端总数
    S_t = 随机选择 m 个客户端
    for 每个客户端 k in S_t (并行):
      w_{t+1}^k = ClientUpdate(k, w_t)
    w_{t+1} = sum (n_k / n) x w_{t+1}^k  // 加权平均

客户端 k 执行:
  B = 把本地数据分成若干批次
  for 本地 epoch e = 1, ..., E:
    for 批次 b in B:
      w = w - lr x grad_loss(w; b)
  return w

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Subset
import numpy as np
from copy import deepcopy
from typing import List, Dict, Tuple
import random


# ========== 模型定义 ==========

class SimpleNet(nn.Module):
    """简单的分类网络"""
    def __init__(self, input_dim: int, hidden_dim: int, num_classes: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, num_classes)
        )

    def forward(self, x):
        return self.net(x)


# ========== 客户端 ==========

class FLClient:
    """联邦学习客户端"""

    def __init__(
        self,
        client_id: int,
        dataset,
        device: str = 'cpu'
    ):
        self.client_id = client_id
        self.dataset = dataset
        self.device = device

    def local_train(
        self,
        model: nn.Module,
        local_epochs: int,
        batch_size: int,
        lr: float
    ) -> Tuple[Dict, float]:
        """
        用本地数据训练模型
        Returns: (更新后的权重, 本地损失)
        """
        model = deepcopy(model).to(self.device)
        model.train()

        loader = DataLoader(
            self.dataset,
            batch_size=batch_size,
            shuffle=True
        )
        optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        n_batches = 0

        for epoch in range(local_epochs):
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)
                optimizer.zero_grad()
                output = model(X)
                loss = criterion(output, y)
                loss.backward()
                optimizer.step()

                total_loss += loss.item()
                n_batches += 1

        avg_loss = total_loss / max(n_batches, 1)
        return model.state_dict(), avg_loss

    def evaluate(self, model: nn.Module) -> Tuple[float, float]:
        """用本地数据评估模型"""
        model = deepcopy(model).to(self.device)
        model.eval()

        loader = DataLoader(self.dataset, batch_size=64)
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)
                output = model(X)
                loss = criterion(output, y)
                total_loss += loss.item()

                _, predicted = output.max(1)
                total += y.size(0)
                correct += predicted.eq(y).sum().item()

        return total_loss / len(loader), correct / total


# ========== 服务器 ==========

class FedAvgServer:
    """FedAvg 服务器"""

    def __init__(
        self,
        global_model: nn.Module,
        clients: List[FLClient],
        fraction: float = 0.1,  # C
        device: str = 'cpu'
    ):
        self.global_model = global_model.to(device)
        self.clients = clients
        self.fraction = fraction
        self.device = device
        self.round_history = []

    def select_clients(self) -> List[FLClient]:
        """在每一轮选择客户端"""
        m = max(int(self.fraction * len(self.clients)), 1)
        return random.sample(self.clients, m)

    def aggregate(
        self,
        client_weights: List[Dict],
        client_sizes: List[int]
    ) -> Dict:
        """
        用加权平均聚合模型(FedAvg)
        按 n_k / n 的比例做加权平均
        """
        total_size = sum(client_sizes)
        aggregated = {}

        for key in client_weights[0].keys():
            aggregated[key] = torch.zeros_like(
                client_weights[0][key], dtype=torch.float32
            )
            for w, size in zip(client_weights, client_sizes):
                weight = size / total_size
                aggregated[key] += weight * w[key].float()

        return aggregated

    def train_round(
        self,
        local_epochs: int = 5,
        batch_size: int = 32,
        lr: float = 0.01
    ) -> Dict:
        """执行一轮 FL"""
        selected = self.select_clients()

        client_weights = []
        client_sizes = []
        client_losses = []

        for client in selected:
            weights, loss = client.local_train(
                self.global_model, local_epochs, batch_size, lr
            )
            client_weights.append(weights)
            client_sizes.append(len(client.dataset))
            client_losses.append(loss)

        # 聚合
        new_weights = self.aggregate(client_weights, client_sizes)
        self.global_model.load_state_dict(new_weights)

        round_info = {
            'num_clients': len(selected),
            'avg_local_loss': np.mean(client_losses),
            'client_losses': client_losses
        }
        self.round_history.append(round_info)
        return round_info

    def evaluate_global(self, test_loader: DataLoader) -> Tuple[float, float]:
        """评估全局模型"""
        self.global_model.eval()
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for X, y in test_loader:
                X, y = X.to(self.device), y.to(self.device)
                output = self.global_model(X)
                loss = criterion(output, y)
                total_loss += loss.item()

                _, predicted = output.max(1)
                total += y.size(0)
                correct += predicted.eq(y).sum().item()

        return total_loss / len(test_loader), correct / total

    def federated_train(
        self,
        num_rounds: int,
        local_epochs: int = 5,
        batch_size: int = 32,
        lr: float = 0.01,
        test_loader: DataLoader = None
    ):
        """完整的 FL 训练循环"""
        print(f"联邦学习开始: {num_rounds} 轮, {len(self.clients)} 个客户端")

        for round_num in range(1, num_rounds + 1):
            round_info = self.train_round(local_epochs, batch_size, lr)

            if test_loader and round_num % 10 == 0:
                test_loss, test_acc = self.evaluate_global(test_loader)
                print(
                    f"轮次 {round_num:3d}/{num_rounds} | "
                    f"客户端: {round_info['num_clients']} | "
                    f"本地损失: {round_info['avg_local_loss']:.4f} | "
                    f"测试准确率: {test_acc:.4f}"
                )

        print("联邦学习完成！")


# ========== 数据分布(Non-IID 模拟) ==========

def create_non_iid_partition(
    dataset,
    num_clients: int,
    num_classes: int,
    alpha: float = 0.5
) -> List[List[int]]:
    """
    使用 Dirichlet 分布进行 Non-IID 数据划分
    alpha 越低，分布越不均衡
    """
    labels = np.array([dataset[i][1] for i in range(len(dataset))])
    client_indices = [[] for _ in range(num_clients)]

    for cls in range(num_classes):
        cls_indices = np.where(labels == cls)[0]
        np.random.shuffle(cls_indices)

        # 用 Dirichlet 分布分配给各客户端
        proportions = np.random.dirichlet(
            [alpha] * num_clients
        )
        proportions = (proportions * len(cls_indices)).astype(int)
        proportions[-1] = len(cls_indices) - proportions[:-1].sum()

        start = 0
        for k, prop in enumerate(proportions):
            client_indices[k].extend(
                cls_indices[start:start + prop].tolist()
            )
            start += prop

    return client_indices


# ========== 主程序 ==========

def run_fedavg_demo():
    """运行 FedAvg 演示"""
    import torchvision
    import torchvision.transforms as transforms

    # 加载 MNIST 数据集
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = torchvision.datasets.MNIST(
        './data', train=True, download=True, transform=transform
    )
    test_dataset = torchvision.datasets.MNIST(
        './data', train=False, transform=transform
    )

    # Non-IID 数据划分
    num_clients = 20
    client_indices = create_non_iid_partition(
        train_dataset, num_clients, num_classes=10, alpha=0.5
    )

    # 创建客户端
    clients = []
    for k in range(num_clients):
        subset = Subset(train_dataset, client_indices[k])
        clients.append(FLClient(k, subset))

    print(f"客户端数量: {num_clients}")
    print(f"每个客户端的平均样本数: {np.mean([len(c.dataset) for c in clients]):.1f}")

    # 创建全局模型
    global_model = SimpleNet(784, 256, 10)

    # 创建服务器(选择 10% 的客户端)
    server = FedAvgServer(global_model, clients, fraction=0.1)

    # 测试数据加载器
    test_loader = DataLoader(test_dataset, batch_size=256)

    # FL 训练
    server.federated_train(
        num_rounds=100,
        local_epochs=5,
        batch_size=32,
        lr=0.01,
        test_loader=test_loader
    )

    # 最终评估
    _, final_acc = server.evaluate_global(test_loader)
    print(f"\n最终测试准确率: {final_acc:.4f}")


if __name__ == '__main__':
    run_fedavg_demo()

h_k(w; w^t) = F_k(w) + (mu/2) x ||w - w^t||^2

class FedProxClient(FLClient):
    """FedProx 客户端：添加近端项"""

    def local_train_prox(
        self,
        model: nn.Module,
        global_weights: Dict,
        local_epochs: int,
        batch_size: int,
        lr: float,
        mu: float = 0.01  # 近端项权重
    ) -> Tuple[Dict, float]:
        """
        包含近端项的本地训练
        h_k(w) = F_k(w) + (mu/2) * ||w - w^t||^2
        """
        model = deepcopy(model).to(self.device)
        model.train()

        # 把全局权重保存为参照点
        global_model = deepcopy(model)
        global_model.load_state_dict(global_weights)
        for param in global_model.parameters():
            param.requires_grad = False

        loader = DataLoader(
            self.dataset,
            batch_size=batch_size,
            shuffle=True
        )
        optimizer = optim.SGD(model.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        n_batches = 0

        for epoch in range(local_epochs):
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)
                optimizer.zero_grad()

                output = model(X)
                task_loss = criterion(output, y)

                # 近端项: (mu/2) * ||w - w_global||^2
                prox_loss = 0.0
                for w, w_global in zip(
                    model.parameters(),
                    global_model.parameters()
                ):
                    prox_loss += (mu / 2) * torch.norm(w - w_global) ** 2

                total_batch_loss = task_loss + prox_loss
                total_batch_loss.backward()
                optimizer.step()

                total_loss += task_loss.item()
                n_batches += 1

        avg_loss = total_loss / max(n_batches, 1)
        return model.state_dict(), avg_loss

class ScaffoldClient:
    """SCAFFOLD 客户端"""

    def __init__(self, client_id, dataset, device='cpu'):
        self.client_id = client_id
        self.dataset = dataset
        self.device = device
        # 客户端控制变量初始化
        self.c_k = None

    def init_control_variate(self, model: nn.Module):
        """初始化控制变量"""
        self.c_k = {
            name: torch.zeros_like(param)
            for name, param in model.named_parameters()
        }

    def local_train_scaffold(
        self,
        model: nn.Module,
        server_control: Dict,
        local_epochs: int,
        batch_size: int,
        lr: float
    ) -> Tuple[Dict, Dict, float]:
        """
        SCAFFOLD 本地训练
        Returns: (更新后的权重, 控制变量更新, 损失)
        """
        if self.c_k is None:
            self.init_control_variate(model)

        model = deepcopy(model).to(self.device)
        model.train()

        initial_weights = deepcopy(model.state_dict())
        loader = DataLoader(self.dataset, batch_size=batch_size, shuffle=True)
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        n_batches = 0
        total_steps = local_epochs * len(loader)

        for epoch in range(local_epochs):
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)

                # 手动更新参数(包含 SCAFFOLD 校正)
                output = model(X)
                loss = criterion(output, y)
                loss.backward()

                with torch.no_grad():
                    for name, param in model.named_parameters():
                        if param.grad is not None:
                            # SCAFFOLD 校正: g - c_k + c
                            correction = (
                                self.c_k[name].to(self.device)
                                - server_control[name].to(self.device)
                            )
                            param -= lr * (param.grad + correction)
                            param.grad.zero_()

                total_loss += loss.item()
                n_batches += 1

        final_weights = model.state_dict()

        # 更新控制变量
        # c_k+ = c_k - c + (1 / (K * lr)) * (w_0 - w_K)
        new_c_k = {}
        c_k_diff = {}
        for name in self.c_k:
            w_diff = (
                initial_weights[name].float() - final_weights[name].float()
            )
            new_c_k[name] = (
                self.c_k[name]
                - server_control[name]
                + w_diff / (total_steps * lr)
            )
            c_k_diff[name] = new_c_k[name] - self.c_k[name]

        self.c_k = new_c_k
        avg_loss = total_loss / max(n_batches, 1)
        return final_weights, c_k_diff, avg_loss

对所有相邻数据集 D、D' 及输出集合 S:
Pr[M(D) ∈ S] ≤ exp(ε) x Pr[M(D') ∈ S] + δ

g_clipped = g x min(1, C / ||g||_2)

g_dp = g_clipped + N(0, σ^2 x C^2 x I)

import torch
import torch.nn as nn
from opacus import PrivacyEngine
from opacus.validators import ModuleValidator


def make_private_model(model: nn.Module) -> nn.Module:
    """转换为 Opacus 兼容模型(BatchNorm -> GroupNorm)"""
    model = ModuleValidator.fix(model)
    return model


class DPFLClient:
    """应用了差分隐私的 FL 客户端"""

    def __init__(
        self,
        client_id: int,
        dataset,
        target_epsilon: float = 1.0,    # 隐私预算
        target_delta: float = 1e-5,     # 失败概率
        max_grad_norm: float = 1.0,     # 梯度裁剪阈值
        device: str = 'cpu'
    ):
        self.client_id = client_id
        self.dataset = dataset
        self.target_epsilon = target_epsilon
        self.target_delta = target_delta
        self.max_grad_norm = max_grad_norm
        self.device = device

    def dp_local_train(
        self,
        model: nn.Module,
        local_epochs: int,
        batch_size: int,
        lr: float
    ) -> Tuple[Dict, float, float]:
        """
        差分隐私本地训练
        Returns: (权重, 损失, 消耗的 epsilon)
        """
        model = make_private_model(deepcopy(model)).to(self.device)
        model.train()

        loader = DataLoader(
            self.dataset,
            batch_size=batch_size,
            shuffle=True,
            drop_last=True  # Opacus 的要求
        )
        optimizer = optim.SGD(model.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()

        # 挂载 Opacus PrivacyEngine
        privacy_engine = PrivacyEngine()
        model, optimizer, loader = privacy_engine.make_private_with_epsilon(
            module=model,
            optimizer=optimizer,
            data_loader=loader,
            epochs=local_epochs,
            target_epsilon=self.target_epsilon,
            target_delta=self.target_delta,
            max_grad_norm=self.max_grad_norm
        )

        total_loss = 0.0
        n_batches = 0

        for epoch in range(local_epochs):
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)
                optimizer.zero_grad()
                output = model(X)
                loss = criterion(output, y)
                loss.backward()
                optimizer.step()

                total_loss += loss.item()
                n_batches += 1

        epsilon_used = privacy_engine.get_epsilon(self.target_delta)
        avg_loss = total_loss / max(n_batches, 1)

        # 去除 Opacus 包装，返回纯净权重
        clean_weights = {
            k.replace('_module.', ''): v
            for k, v in model.state_dict().items()
        }

        return clean_weights, avg_loss, epsilon_used


class DPFLServer:
    """DP-FL 服务器: 服务器端 DP 聚合"""

    def __init__(
        self,
        global_model: nn.Module,
        clients: List['DPFLClient'],
        noise_multiplier: float = 0.5,
        max_grad_norm: float = 1.0,
        device: str = 'cpu'
    ):
        self.global_model = global_model.to(device)
        self.clients = clients
        self.noise_multiplier = noise_multiplier
        self.max_grad_norm = max_grad_norm
        self.device = device

    def clip_and_aggregate(
        self,
        client_weights: List[Dict],
        reference_weights: Dict
    ) -> Dict:
        """
        服务器端裁剪与加噪
        先裁剪每个更新，再加入高斯噪声
        """
        n_clients = len(client_weights)

        # 计算更新量(delta_w = w_k - w_global)
        updates = []
        for w in client_weights:
            delta = {
                k: w[k].float() - reference_weights[k].float()
                for k in reference_weights
            }
            updates.append(delta)

        # L2 范数裁剪
        clipped_updates = []
        for delta in updates:
            total_norm = torch.sqrt(
                sum(torch.norm(v) ** 2 for v in delta.values())
            )
            clip_factor = min(1.0, self.max_grad_norm / (total_norm + 1e-8))
            clipped = {k: v * clip_factor for k, v in delta.items()}
            clipped_updates.append(clipped)

        # 求和
        summed = {}
        for key in reference_weights:
            summed[key] = sum(u[key] for u in clipped_updates)

        # 加入高斯噪声
        sigma = self.noise_multiplier * self.max_grad_norm
        noisy = {}
        for key in summed:
            noise = torch.randn_like(summed[key]) * sigma
            noisy[key] = (summed[key] + noise) / n_clients

        # 应用到全局模型
        aggregated = {
            k: reference_weights[k].float() + noisy[k]
            for k in reference_weights
        }
        return aggregated

import secrets
import hashlib


class SecureAggregation:
    """
    简化版的安全聚合模拟
    实际实现会使用 Diffie-Hellman 密钥交换
    """

    def __init__(self, num_clients: int, seed_length: int = 32):
        self.num_clients = num_clients
        self.seed_length = seed_length

    def generate_pairwise_masks(
        self,
        client_id: int,
        model_shape: Dict[str, torch.Size],
        shared_seeds: Dict
    ) -> Dict:
        """
        生成客户端两两之间的掩码
        在 client i 与 j 之间: i>j 则相加, i<j 则相减
        """
        mask = {k: torch.zeros(v) for k, v in model_shape.items()}

        for other_id, seed in shared_seeds.items():
            # 用共享种子生成伪随机掩码
            torch.manual_seed(seed)
            direction = 1 if client_id > other_id else -1

            for key in mask:
                mask[key] += direction * torch.randn(model_shape[key])

        return mask

    def client_mask_update(
        self,
        local_update: Dict,
        mask: Dict
    ) -> Dict:
        """把掩码应用到更新上"""
        return {k: local_update[k] + mask[k] for k in local_update}

    def server_aggregate_masked(
        self,
        masked_updates: List[Dict]
    ) -> Dict:
        """对带掩码的更新求和(掩码会相互抵消)"""
        aggregated = {}
        for key in masked_updates[0]:
            aggregated[key] = sum(u[key] for u in masked_updates)
            aggregated[key] /= len(masked_updates)
        return aggregated

pip install flwr
pip install flwr[simulation]

import flwr as fl
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from typing import Dict, List, Tuple
import numpy as np


class FlowerClient(fl.client.NumPyClient):
    """Flower FL 客户端"""

    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        device: str = 'cpu'
    ):
        self.model = model.to(device)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = device
        self.criterion = nn.CrossEntropyLoss()

    def get_parameters(self, config: Dict) -> List[np.ndarray]:
        """把当前模型参数以 NumPy 数组形式返回"""
        return [
            val.cpu().numpy()
            for _, val in self.model.state_dict().items()
        ]

    def set_parameters(self, parameters: List[np.ndarray]):
        """用 NumPy 数组设置模型参数"""
        params_dict = zip(self.model.state_dict().keys(), parameters)
        state_dict = {
            k: torch.tensor(v)
            for k, v in params_dict
        }
        self.model.load_state_dict(state_dict, strict=True)

    def fit(
        self,
        parameters: List[np.ndarray],
        config: Dict
    ) -> Tuple[List[np.ndarray], int, Dict]:
        """用服务器发来的参数进行本地训练"""
        # 设置全局参数
        self.set_parameters(parameters)

        # 解析配置参数
        local_epochs = int(config.get('local_epochs', 1))
        lr = float(config.get('lr', 0.01))

        # 本地训练
        self.model.train()
        optimizer = torch.optim.SGD(
            self.model.parameters(), lr=lr, momentum=0.9
        )

        train_loss = 0.0
        for epoch in range(local_epochs):
            for X, y in self.train_loader:
                X, y = X.to(self.device), y.to(self.device)
                optimizer.zero_grad()
                loss = self.criterion(self.model(X), y)
                loss.backward()
                optimizer.step()
                train_loss += loss.item()

        return (
            self.get_parameters(config={}),
            len(self.train_loader.dataset),
            {'train_loss': train_loss}
        )

    def evaluate(
        self,
        parameters: List[np.ndarray],
        config: Dict
    ) -> Tuple[float, int, Dict]:
        """用服务器发来的参数进行评估"""
        self.set_parameters(parameters)
        self.model.eval()

        val_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for X, y in self.val_loader:
                X, y = X.to(self.device), y.to(self.device)
                output = self.model(X)
                val_loss += self.criterion(output, y).item()

                _, predicted = output.max(1)
                total += y.size(0)
                correct += predicted.eq(y).sum().item()

        accuracy = correct / total
        return (
            val_loss / len(self.val_loader),
            len(self.val_loader.dataset),
            {'accuracy': accuracy}
        )

from flwr.server.strategy import FedAvg
from flwr.common import Parameters, FitIns, FitRes, EvaluateIns, EvaluateRes
from flwr.server.client_proxy import ClientProxy


class CustomFedAvgStrategy(FedAvg):
    """自定义 FedAvg 策略"""

    def __init__(
        self,
        fraction_fit: float = 0.1,
        fraction_evaluate: float = 0.1,
        min_fit_clients: int = 2,
        min_evaluate_clients: int = 2,
        min_available_clients: int = 2,
        initial_parameters=None,
    ):
        super().__init__(
            fraction_fit=fraction_fit,
            fraction_evaluate=fraction_evaluate,
            min_fit_clients=min_fit_clients,
            min_evaluate_clients=min_evaluate_clients,
            min_available_clients=min_available_clients,
            initial_parameters=initial_parameters,
        )
        self.round_metrics = []

    def configure_fit(
        self,
        server_round: int,
        parameters: Parameters,
        client_manager
    ) -> List[Tuple[ClientProxy, FitIns]]:
        """配置每一轮的训练"""
        # 轮次越多，学习率越低
        lr = 0.01 * (0.99 ** server_round)

        config = {
            'local_epochs': 5,
            'lr': lr,
            'round': server_round,
        }

        fit_ins = FitIns(parameters, config)
        clients = client_manager.sample(
            num_clients=max(
                int(client_manager.num_available() * self.fraction_fit), 1
            ),
            min_num_clients=self.min_fit_clients
        )
        return [(client, fit_ins) for client in clients]

    def aggregate_evaluate(
        self,
        server_round: int,
        results: List[Tuple[ClientProxy, EvaluateRes]],
        failures
    ):
        """聚合并记录评估结果"""
        aggregated_loss, metrics = super().aggregate_evaluate(
            server_round, results, failures
        )

        if metrics:
            print(f"轮次 {server_round}: 聚合准确率 = {metrics.get('accuracy', 0):.4f}")
            self.round_metrics.append({
                'round': server_round,
                'loss': aggregated_loss,
                'metrics': metrics
            })

        return aggregated_loss, metrics


def run_flower_simulation():
    """运行 Flower 模拟"""
    import torchvision
    import torchvision.transforms as transforms
    from torch.utils.data import Subset

    # 准备数据
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = torchvision.datasets.MNIST(
        './data', train=True, download=True, transform=transform
    )
    test_dataset = torchvision.datasets.MNIST(
        './data', train=False, transform=transform
    )

    num_clients = 10

    def client_fn(cid: str):
        """客户端创建函数"""
        cid = int(cid)

        # 数据划分
        n = len(train_dataset) // num_clients
        start = cid * n
        indices = list(range(start, min(start + n, len(train_dataset))))
        train_subset = Subset(train_dataset, indices)

        val_indices = list(range(cid * 100, (cid + 1) * 100))
        val_subset = Subset(test_dataset, val_indices)

        model = SimpleNet(784, 256, 10)

        return FlowerClient(
            model=model,
            train_loader=DataLoader(train_subset, batch_size=32),
            val_loader=DataLoader(val_subset, batch_size=64)
        ).to_client()

    # 初始化全局模型参数
    global_model = SimpleNet(784, 256, 10)
    initial_params = fl.common.ndarrays_to_parameters(
        [val.numpy() for val in global_model.state_dict().values()]
    )

    # 配置策略
    strategy = CustomFedAvgStrategy(
        fraction_fit=0.5,
        fraction_evaluate=0.5,
        min_fit_clients=2,
        min_evaluate_clients=2,
        min_available_clients=num_clients,
        initial_parameters=initial_params
    )

    # 运行模拟
    history = fl.simulation.start_simulation(
        client_fn=client_fn,
        num_clients=num_clients,
        config=fl.server.ServerConfig(num_rounds=50),
        strategy=strategy,
    )

    print("\n=== Flower 模拟完成 ===")
    print(f"最终分布式损失: {history.losses_distributed[-1]}")

import torch
import torch.nn as nn
import torchvision.models as models
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os
import pandas as pd


class ChestXRayDataset(Dataset):
    """胸部 X 光数据集(按医院划分)"""

    def __init__(
        self,
        data_dir: str,
        labels_file: str,
        transform=None,
        hospital_id: int = 0
    ):
        self.data_dir = data_dir
        self.labels = pd.read_csv(labels_file)
        self.transform = transform
        self.hospital_id = hospital_id

        # 只筛选该医院的数据
        self.labels = self.labels[
            self.labels['hospital_id'] == hospital_id
        ].reset_index(drop=True)

        self.classes = [
            'Normal', 'Pneumonia', 'COVID-19', 'Tuberculosis'
        ]

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        img_path = os.path.join(
            self.data_dir, self.labels.iloc[idx]['filename']
        )
        image = Image.open(img_path).convert('RGB')
        label = self.labels.iloc[idx]['label']

        if self.transform:
            image = self.transform(image)

        return image, label


def build_chest_model(num_classes: int = 4, pretrained: bool = True):
    """基于 ResNet50 的胸部 X 光分类模型"""
    model = models.resnet50(pretrained=pretrained)

    # 替换最后一层
    in_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Dropout(0.3),
        nn.Linear(in_features, 256),
        nn.ReLU(),
        nn.Linear(256, num_classes)
    )
    return model


class HospitalFLClient(FlowerClient):
    """面向医院的 FL 客户端(处理类别不均衡)"""

    def __init__(
        self,
        hospital_id: int,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        class_weights: torch.Tensor = None,
        device: str = 'cpu'
    ):
        super().__init__(model, train_loader, val_loader, device)
        self.hospital_id = hospital_id

        if class_weights is not None:
            self.criterion = nn.CrossEntropyLoss(
                weight=class_weights.to(device)
            )
        else:
            self.criterion = nn.CrossEntropyLoss()

    def compute_class_weights(self, dataset) -> torch.Tensor:
        """计算用于校正类别不均衡的权重"""
        labels = [dataset[i][1] for i in range(len(dataset))]
        class_counts = torch.bincount(torch.tensor(labels))
        weights = 1.0 / class_counts.float()
        weights = weights / weights.sum()
        return weights


def setup_hospital_federation(num_hospitals: int = 5):
    """
    配置多医院联邦学习
    实际场景中，每家医院的数据都保存在本地
    """
    import torchvision.transforms as transforms

    # 数据预处理
    train_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(10),
        transforms.ColorJitter(brightness=0.2, contrast=0.2),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])

    val_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])

    hospitals = []
    for h_id in range(num_hospitals):
        # 实际场景中会在各医院自己的服务器上运行
        # 这里为了模拟，统一放在同一台服务器上运行
        print(f"正在准备医院 {h_id} 的数据...")

        # 生成占位数据(实际实现时应替换为真实数据)
        # 这里只展示示例结构

    return hospitals


# 联邦学习训练循环
def train_hospital_federation():
    """运行医院联邦学习"""
    print("=== 多医院联邦学习开始 ===")
    print("各医院的患者数据不会被发送到外部。")
    print("只会聚合模型参数的更新。")

    global_model = build_chest_model(num_classes=4)

    print(f"\n全局模型参数数量: {sum(p.numel() for p in global_model.parameters()):,}")
    print(f"通信成本(FP32): {sum(p.numel() for p in global_model.parameters()) * 4 / 1024 / 1024:.1f} MB/轮")