联邦学习详解¶

1. 服务器下发全局模型参数
2. 各客户端（医院）用本地数据训练
3. 各客户端上传模型梯度/参数更新
4. 服务器聚合更新（如 FedAvg 加权平均）
5. 更新全局模型，重复

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import copy
import numpy as np

# ===== 模拟联邦学习：3 家医院共同训练疾病分类模型 =====

# ===== 1. 定义简单的分类模型 =====
class DiagnosisModel(nn.Module):
    """
    简单的疾病分类神经网络
    输入：基因表达特征
    输出：疾病类别概率
    """
    def __init__(self, input_dim=50, hidden_dim=32, num_classes=2):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),   # 特征提取层
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, num_classes)  # 分类输出层
        )

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

# ===== 2. 模拟各医院的本地数据（非 IID，各医院患者群体不同）=====
def generate_hospital_data(hospital_id, n_samples=200, n_features=50):
    """
    为每个医院生成模拟数据
    医院 ID 影响数据分布（模拟非 IID 情况）
    """
    np.random.seed(hospital_id * 42)

    # 各医院有不同的基因表达背景（模拟不同患者群体）
    X = np.random.randn(n_samples, n_features).astype(np.float32)
    X += hospital_id * 0.3  # 各医院数据有系统性偏移（批次效应）

    # 标签：后20个特征和相关
    y = (X[:, 40:50].sum(axis=1) > 0).astype(np.int64)

    X_tensor = torch.FloatTensor(X)
    y_tensor = torch.LongTensor(y)
    dataset = TensorDataset(X_tensor, y_tensor)
    return DataLoader(dataset, batch_size=32, shuffle=True)

# 创建 3 个医院的数据
num_hospitals = 3
hospital_dataloaders = [generate_hospital_data(i) for i in range(num_hospitals)]

# ===== 3. FedAvg 算法实现 =====
def local_train(model, dataloader, epochs=5, lr=0.01):
    """
    本地训练：各医院在本地数据上训练几轮
    返回更新后的模型参数
    """
    model.train()
    optimizer = optim.SGD(model.parameters(), lr=lr)  # 随机梯度下降
    loss_fn = nn.CrossEntropyLoss()                    # 交叉熵损失

    for epoch in range(epochs):
        for X_batch, y_batch in dataloader:
            optimizer.zero_grad()
            output = model(X_batch)            # 前向传播
            loss = loss_fn(output, y_batch)    # 计算损失
            loss.backward()                    # 反向传播
            optimizer.step()                   # 更新本地参数

    return model.state_dict()  # 返回更新后的模型参数字典

def fedavg_aggregate(global_model, local_state_dicts, weights=None):
    """
    FedAvg：按权重平均各医院的模型参数
    weights: 各医院的数据量权重（None=均等）
    """
    if weights is None:
        weights = [1.0 / len(local_state_dicts)] * len(local_state_dicts)

    global_state = global_model.state_dict()

    # 对每一层参数做加权平均
    for key in global_state.keys():
        # 初始化为零张量
        aggregated_param = torch.zeros_like(global_state[key], dtype=torch.float)
        for weight, local_state in zip(weights, local_state_dicts):
            # 累积加权参数（注意转换为 float 避免整型问题）
            aggregated_param += weight * local_state[key].float()
        global_state[key] = aggregated_param.to(global_state[key].dtype)

    global_model.load_state_dict(global_state)  # 更新全局模型
    return global_model

# ===== 4. 联邦训练主循环 =====
global_model = DiagnosisModel(input_dim=50)  # 初始化全局模型
num_rounds = 20  # 联邦学习轮数

print('=== 开始联邦学习训练 ===')
for round_num in range(num_rounds):
    local_states = []  # 收集各医院的本地模型参数

    for hospital_id, dataloader in enumerate(hospital_dataloaders):
        # 每家医院拿到全局模型副本，在本地训练
        local_model = copy.deepcopy(global_model)   # 深拷贝全局模型
        local_state = local_train(local_model, dataloader, epochs=3)
        local_states.append(local_state)             # 收集本地参数更新

    # 服务器聚合（等权重 FedAvg）
    global_model = fedavg_aggregate(global_model, local_states)

    if round_num % 5 == 0:
        # 简单评估全局模型（实际应在验证集上）
        print(f'Round {round_num}: 联邦聚合完成')

print('\n联邦学习完成！全局模型已训练完毕。')
print('注意：整个过程中，各医院的原始数据从未离开本地。')