06-联邦学习¶

def top_k_sparsification(gradient, k_ratio=0.01):
    """Top-K梯度稀疏化"""
    flat_grad = gradient.flatten()
    k = max(1, int(flat_grad.numel() * k_ratio))

    # 取绝对值最大的Top-K个元素
    values, indices = flat_grad.abs().topk(k)
    sparse_grad = torch.zeros_like(flat_grad)
    sparse_grad[indices] = flat_grad[indices]

    return sparse_grad.view_as(gradient)

def quantize_gradient(gradient, n_bits=8):
    """低精度梯度量化"""
    min_val = gradient.min()
    max_val = gradient.max()
    n_levels = 2 ** n_bits - 1

    # 线性量化
    scale = (max_val - min_val) / n_levels
    quantized = torch.round((gradient - min_val) / scale).to(torch.uint8)

    # 解量化
    dequantized = quantized.float() * scale + min_val
    return dequantized

class FedPerModel(nn.Module):  # 继承nn.Module定义神经网络层
    def __init__(self):  # __init__构造方法，创建对象时自动调用
        super().__init__()  # super()调用父类方法
        # 共享层（参与联邦聚合）
        self.shared_layers = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        # 个性化层（仅本地训练）
        self.personal_layers = nn.Sequential(
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 10),
        )

    def forward(self, x):
        x = self.shared_layers(x).flatten(1)
        return self.personal_layers(x)

    def get_shared_params(self):
        return dict(self.shared_layers.named_parameters())

    def get_personal_params(self):
        return dict(self.personal_layers.named_parameters())

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Subset
import numpy as np
from copy import deepcopy

# ==================== 模型定义 ====================
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 32 * 7 * 7)  # view重塑张量形状（要求内存连续）
        x = F.relu(self.fc1(x))
        return self.fc2(x)

# ==================== 数据划分（模拟Non-IID） ====================
def create_non_iid_partitions(dataset, n_clients, n_shards_per_client=2):
    """按标签排序后切片，模拟Label-skew Non-IID"""
    labels = np.array([dataset[i][1] for i in range(len(dataset))])  # 列表推导式，简洁创建列表  # np.array创建NumPy数组
    sorted_indices = np.argsort(labels)

    n_shards = n_clients * n_shards_per_client
    shard_size = len(dataset) // n_shards
    shards = [sorted_indices[i * shard_size:(i + 1) * shard_size]
              for i in range(n_shards)]

    np.random.shuffle(shards)
    client_indices = {}
    for i in range(n_clients):
        client_indices[i] = np.concatenate(
            shards[i * n_shards_per_client:(i + 1) * n_shards_per_client]
        )
    return client_indices

# ==================== 客户端本地训练 ====================
def client_update(model, dataloader, lr=0.01, local_epochs=5):
    """客户端本地训练"""
    model.train()  # train()开启训练模式
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    for epoch in range(local_epochs):
        for x, y in dataloader:
            optimizer.zero_grad()  # 清零梯度，防止梯度累积
            loss = F.cross_entropy(model(x), y)
            loss.backward()  # 反向传播计算梯度
            optimizer.step()  # 根据梯度更新模型参数

    return model.state_dict()

# ==================== 服务器聚合 ====================
def fedavg_aggregate(global_model, client_state_dicts, client_sizes):
    """FedAvg聚合：按样本数加权平均"""
    total_size = sum(client_sizes)
    global_state = global_model.state_dict()

    for key in global_state:
        global_state[key] = sum(
            client_state_dicts[i][key].float() * (client_sizes[i] / total_size)
            for i in range(len(client_state_dicts))
        )

    global_model.load_state_dict(global_state)
    return global_model

# ==================== 联邦训练主循环 ====================
def federated_train(train_dataset, test_dataset, n_clients=10,
                    n_rounds=20, local_epochs=5, participation_rate=0.3):
    """完整的联邦训练流程"""
    # 数据划分
    client_indices = create_non_iid_partitions(train_dataset, n_clients)

    # 初始化全局模型
    global_model = SimpleCNN()

    # 测试集loader
    test_loader = DataLoader(test_dataset, batch_size=128)

    for round_idx in range(n_rounds):
        # 1. 选择参与客户端
        n_selected = max(1, int(n_clients * participation_rate))
        selected = np.random.choice(n_clients, n_selected, replace=False)

        client_states = []
        client_sizes = []

        for client_id in selected:
            # 2. 下发全局模型
            local_model = deepcopy(global_model)

            # 3. 本地训练
            indices = client_indices[client_id]
            local_dataset = Subset(train_dataset, indices)
            local_loader = DataLoader(local_dataset, batch_size=32, shuffle=True)

            state_dict = client_update(local_model, local_loader,
                                       local_epochs=local_epochs)
            client_states.append(state_dict)
            client_sizes.append(len(indices))

        # 4. 服务器聚合
        global_model = fedavg_aggregate(global_model, client_states, client_sizes)

        # 5. 评估
        if (round_idx + 1) % 5 == 0:
            global_model.eval()  # eval()开启评估模式（关闭Dropout等）
            correct, total = 0, 0
            with torch.no_grad():  # 禁用梯度计算，节省内存
                for x, y in test_loader:
                    pred = global_model(x).argmax(dim=1)
                    correct += (pred == y).sum().item()  # .item()将单元素张量转为Python数值
                    total += y.size(0)
            print(f"Round {round_idx+1}, Accuracy: {correct/total:.4f}")

    return global_model

# 使用示例（需要torchvision的MNIST数据集）
# from torchvision import datasets, transforms
# train_set = datasets.MNIST('./data', train=True, download=True,
#                            transform=transforms.ToTensor())
# test_set = datasets.MNIST('./data', train=False,
#                           transform=transforms.ToTensor())
# model = federated_train(train_set, test_set)