联邦学习（一）：隐私保护的分布式训练

# 成员推断攻击示例：判断样本是否在训练集中
import numpy as np
from sklearn.metrics import roc_auc_score

def membership_inference_attack(model, shadow_train, shadow_test, target_samples):
    """使用 shadow model 模拟成员推断攻击"""
    train_confidence = np.max(model.predict_proba(shadow_train), axis=1)
    test_confidence = np.max(model.predict_proba(shadow_test), axis=1)
    target_confidence = np.max(model.predict_proba(target_samples), axis=1)
    labels = np.concatenate([np.ones(len(train_confidence)),
                             np.zeros(len(test_confidence))])
    scores = np.concatenate([train_confidence, test_confidence])
    auc = roc_auc_score(labels, scores)
    threshold = np.median(scores)
    predictions = target_confidence > threshold
    return predictions, auc

print(f"攻击 AUC: {auc:.4f}，AUC 越高说明成员推断越容易")

# 模型反演攻击：从输出重构输入特征
import torch
import torch.nn as nn

class ModelInversionAttack:
    def __init__(self, model, target_class, input_shape, lr=0.01):
        self.model = model
        self.target_class = target_class
        # 初始化为随机噪声
        self.fake_input = nn.Parameter(torch.randn(input_shape, requires_grad=True))
        self.optimizer = torch.optim.Adam([self.fake_input], lr=lr)

    def reconstruct(self, steps=1000):
        target = torch.tensor([self.target_class], dtype=torch.long)
        for step in range(steps):
            self.optimizer.zero_grad()
            output = self.model(self.fake_input)
            loss = nn.CrossEntropyLoss()(output, target)
            loss.backward()
            self.optimizer.step()
        return self.fake_input.detach()

# 攻击成功后，fake_input 会逼近目标类别的典型训练样本

# 拉普拉斯机制实现差分隐私
import numpy as np

def laplace_mechanism(query_result, sensitivity, epsilon):
    """对查询结果添加拉普拉斯噪声"""
    scale = sensitivity / epsilon
    noise = np.random.laplace(loc=0, scale=scale)
    return query_result + noise

# 示例：统计平均收入（敏感度 = 数据范围 / n）
def private_mean(data, epsilon, data_min=0, data_max=500000):
    n = len(data)
    sensitivity = (data_max - data_min) / n
    true_mean = np.mean(data)
    return laplace_mechanism(true_mean, sensitivity, epsilon)

# ε=1.0 时，单次查询的隐私保障为 e^1 ≈ 2.718 倍概率比
print(f"真实均值: {true_mean:.2f}, 差分隐私均值: {private_mean(data, epsilon=1.0):.2f}")

# 指数机制：从候选集中隐私保护地选择最优项
import numpy as np

def exponential_mechanism(candidates, utility_fn, sensitivity, epsilon):
    """指数机制：按效用函数得分采样候选"""
    scores = np.array([utility_fn(c) for c in candidates])
    # 计算每个候选的指数权重
    alpha = epsilon / (2 * sensitivity)
    weights = np.exp(alpha * scores)
    weights /= weights.sum()
    idx = np.random.choice(len(candidates), p=weights)
    return candidates[idx]

# 示例：隐私保护地选择最佳超参数
param_candidates = [0.001, 0.01, 0.1, 0.5]
def accuracy_score(lr):
    # 模拟不同学习率的验证集准确率
    return 0.85 - abs(lr - 0.05)

best_lr = exponential_mechanism(param_candidates, accuracy_score,
                                 sensitivity=0.01, epsilon=1.0)
print(f"差分隐私选择的超参数: {best_lr}")

# DP-SGD 手动实现（简化版）
import torch
import torch.nn as nn

def dp_sgd_step(model, data_batch, labels, optimizer, noise_multiplier, max_grad_norm):
    """单步 DP-SGD 训练"""
    optimizer.zero_grad()
    # 计算每个样本的梯度（per-sample gradients）
    per_sample_grads = []
    for x, y in zip(data_batch, labels):
        loss = model(x.unsqueeze(0), y.unsqueeze(0))
        loss.backward()
        grad = []
        for p in model.parameters():
            grad.append(p.grad.clone().detach())
            p.grad = None  # 清除累积梯度
        per_sample_grads.append(grad)

    # 梯度裁剪（per-sample clipping）
    clipped_grads = []
    for sample_grad in per_sample_grads:
        total_norm = torch.sqrt(sum((g ** 2).sum() for g in sample_grad))
        clip_coeff = min(1.0, max_grad_norm / (total_norm + 1e-6))
        clipped = [g * clip_coeff for g in sample_grad]
        clipped_grads.append(clipped)

    # 聚合 + 加噪
    batch_size = len(data_batch)
    for i, param in enumerate(model.parameters()):
        avg_grad = torch.stack([cg[i] for cg in clipped_grads]).mean(dim=0)
        noise = torch.normal(0, noise_multiplier * max_grad_norm, avg_grad.shape)
        param.grad = avg_grad + noise

    optimizer.step()

# 使用 Opacus 库实现 DP-SGD
from opacus import PrivacyEngine
from opacus.validators import ModuleValidator
import torch.nn as nn

# 假设已有 model, optimizer, data_loader
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)
model = ModuleValidator.fix(model)  # 修复不兼容层

optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
privacy_engine = PrivacyEngine()

# 绑定隐私引擎到 model 和 optimizer
model, optimizer, data_loader = privacy_engine.make_private(
    module=model,
    optimizer=optimizer,
    data_loader=data_loader,
    noise_multiplier=1.0,      # 噪声乘数 σ
    max_grad_norm=1.0,         # 梯度裁剪阈值
)

# 正常训练循环即可，Opacus 自动处理 per-sample 梯度和加噪
# 训练完成后查看隐私支出
epsilon = privacy_engine.get_epsilon(delta=1e-5)
print(f"隐私支出: ε={epsilon:.2f} (δ=1e-5)")

# FedAvg 算法核心逻辑（简化版）
import copy
import numpy as np
import torch

def fedavg(global_model, client_datasets, client_weights, local_epochs=5, lr=0.01):
    """联邦平均算法"""
    num_clients = len(client_datasets)
    local_updates = []

    for i in range(num_clients):
        # 克隆全局模型到客户端
        local_model = copy.deepcopy(global_model)
        optimizer = torch.optim.SGD(local_model.parameters(), lr=lr)

        # 本地训练
        for epoch in range(local_epochs):
            for x_batch, y_batch in client_datasets[i]:
                optimizer.zero_grad()
                loss = local_model(x_batch, y_batch)
                loss.backward()
                optimizer.step()

        # 计算模型更新量
        update = {}
        for name, param in local_model.named_parameters():
            global_param = dict(global_model.named_parameters())[name]
            update[name] = param.data - global_param.data
        local_updates.append(update)

    # 加权平均聚合
    total_weight = sum(client_weights)
    aggregated = {}
    for name in global_model.state_dict():
        aggregated[name] = sum(w * local_updates[i][name]
                               for i, w in enumerate(client_weights)) / total_weight

    # 更新全局模型
    for name, param in global_model.named_parameters():
        param.data += aggregated[name]
    return global_model

# 使用 Flower 框架实现联邦学习
import flwr as fl
import torch
import torch.nn as nn

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

    def get_parameters(self, config):
        return [val.cpu().numpy() for _, val in self.model.state_dict().items()]

    def set_parameters(self, parameters):
        state_dict = dict(zip(self.model.state_dict().keys(),
                              [torch.tensor(p) for p in parameters]))
        self.model.load_state_dict(state_dict, strict=True)

    def fit(self, parameters, config):
        self.set_parameters(parameters)
        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)
        self.model.train()
        for x, y in self.train_loader:
            optimizer.zero_grad()
            loss = nn.CrossEntropyLoss()(self.model(x), y)
            loss.backward()
            optimizer.step()
        return self.get_parameters(config), len(self.train_loader), {}

    def evaluate(self, parameters, config):
        self.set_parameters(parameters)
        correct, total = 0, 0
        self.model.eval()
        with torch.no_grad():
            for x, y in self.val_loader:
                correct += (self.model(x).argmax(1) == y).sum().item()
                total += len(y)
        return float(correct / total), total, {"accuracy": correct / total}

# 启动客户端
fl.client.start_numpy_client(server_address="0.0.0.0:8080",
                              client=FlowerClient(model, train_loader, val_loader))

# 加法秘密分享（Additive Secret Sharing）示意
import numpy as np

class AdditiveSecretSharing:
    """基于模运算的加法秘密分享（2-out-of-2 方案）"""
    def __init__(self, prime=2**32 - 5):
        self.prime = prime  # 大素数

    def share(self, value, n_parties=2):
        """将 value 拆分为 n_parties 个份额"""
        shares = np.random.randint(0, self.prime, size=n_parties - 1)
        last_share = (value - sum(shares)) % self.prime
        shares = list(shares) + [last_share]
        return shares

    def reconstruct(self, shares):
        """从所有份额恢复原始值"""
        return sum(shares) % self.prime

# 示例：两方安全加法
ss = AdditiveSecretSharing()
x_shares = ss.share(42)
y_shares = ss.share(17)
# 各方在本地计算份额之和（无需交换原始值）
z_shares = [(x_shares[i] + y_shares[i]) % ss.prime for i in range(2)]
result = ss.reconstruct(z_shares)
print(f"安全加法结果: {result}")  # 输出 59

# 使用 CrypTen 进行隐私保护推理
import crypten
import torch

# 初始化 CrypTen（多进程环境）
crypten.init()

# 以加密张量加载模型和数据
@crypten.compiler.compile()
class PrivateModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(10, 5)
        self.fc2 = torch.nn.Linear(5, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        return self.fc2(x)

# 加密推理
def private_inference(model, private_input):
    """在加密状态下执行模型推理"""
    # cryptensor 表示加密张量
    encrypted_input = crypten.cryptensor(private_input)
    encrypted_output = model(encrypted_input)
    # 只有授权方可以解密
    return encrypted_output.get_plain_text()

# 注意：CrypTen 需要多进程设置，此处为示意代码
print("CrypTen 支持在加密状态下执行完整的神经网络推理")

# 使用 Pyfhel 进行同态加密计算
from Pyfhel import Pyfhel

# 初始化同态加密上下文
he = Pyfhel()
he.contextGen(p=65537, m=2**15, sec=128)  # 参数设置
he.keyGen()  # 密钥生成
he.relinKeyGen()  # 重线性化密钥（用于乘法）

# 加密数据
a = he.encryptInt(42)
b = he.encryptInt(17)

# 在加密状态下计算
c_add = a + b       # 同态加法
c_mul = a * b       # 同态乘法

# 解密验证
print(f"解密加法: {he.decryptInt(c_add)}")  # 59
print(f"解密乘法: {he.decryptInt(c_mul)}")  # 714

# CKKS 方案支持浮点数同态运算（适用于 ML）
he_ckks = Pyfhel()
he_ckks.contextGen(p=None, m=215, sec=128, scale=240)
he_ckks.keyGen()
he_ckks.relinKeyGen()

x = he_ckks.encryptFrac([0.5, 1.2, -0.3])
y = he_ckks.encryptFrac([0.1, 0.8, 0.5])
z = x + y
print(f"CKKS 解密: {he_ckks.decryptFrac(z)}")  # [0.6, 2.0, 0.2]

# TenSEAL: 用于机器学习的同态加密库
import tenseal as ts

# 创建 CKKS 上下文
context = ts.context(
    ts.SCHEME_TYPE.CKKS,
    poly_modulus_degree=8192,
    coeff_mod_bit_sizes=[60, 40, 40, 60]
)
context.global_scale = 2**40
context.generate_galois_keys()
context.generate_relin_keys()

# 构建简单的隐私保护推理
class HomomorphicInference:
    def __init__(self, context, weights, bias):
        self.context = context
        self.weights = ts.ckks_vector(context, weights)
        self.bias = ts.ckks_vector(context, bias)

    def linear_layer(self, encrypted_input):
        """同态线性层：y = xW + b"""
        result = encrypted_input.mm(self.weights)
        result = result + self.bias
        return result

# 使用示例
weights = [[0.1, 0.2], [0.3, 0.4]]
bias = [0.5, 0.6]
model = HomomorphicInference(context, weights, bias)

client_input = ts.ckks_vector(context, [1.0, 2.0])
encrypted_result = model.linear_layer(client_input)
# 客户端本地解密
print(f"推理结果: {encrypted_result.decrypt()}")

# PySyft: 远程执行 + 差分隐私
import syft as sy
import torch

# 创建虚拟工作节点
alice = sy.VirtualWorker(hook=None, id="alice")
bob = sy.VirtualWorker(hook=None, id="bob")

# 发送数据到远程节点（数据不出本地机器，只在虚拟节点间传递引用）
x = torch.tensor([1, 2, 3, 4]).send(alice)
y = torch.tensor([5, 6, 7, 8]).send(bob)

# 远程计算：结果仍留在远程节点
z = x + y.send(alice)

# 差分隐私训练：使用 Opacus 集成
from opacus import PrivacyEngine
import torch.nn as nn

model = nn.Linear(10, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
privacy_engine = PrivacyEngine()
model, optimizer, dataloader = privacy_engine.make_private(
    module=model, optimizer=optimizer, data_loader=train_loader,
    noise_multiplier=0.8, max_grad_norm=0.5
)
# 训练循环与常规 PyTorch 一致，Opacus 自动注入隐私保护
print("PySyft + Opacus: 远程数据 + 差分隐私训练")

# Flower: 带差分隐私的联邦学习
import flwr as fl
import torch
from opacus import PrivacyEngine

class DPFedAvgClient(fl.client.NumPyClient):
    """集成差分隐私的 Flower 客户端"""
    def __init__(self, model, train_loader, val_loader,
                 noise_multiplier=1.0, max_grad_norm=1.0):
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.privacy_engine = PrivacyEngine()
        # 包装 optimizer 实现 DP-SGD
        self.optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
        self.model, self.optimizer, _ = self.privacy_engine.make_private(
            module=self.model, optimizer=self.optimizer,
            data_loader=self.train_loader,
            noise_multiplier=noise_multiplier,
            max_grad_norm=max_grad_norm,
        )

    def fit(self, parameters, config):
        self.set_parameters(parameters)
        self.model.train()
        for epoch in range(3):
            for x, y in self.train_loader:
                self.optimizer.zero_grad()
                loss = torch.nn.CrossEntropyLoss()(self.model(x), y)
                loss.backward()
                self.optimizer.step()
        epsilon = self.privacy_engine.get_epsilon(delta=1e-5)
        return self.get_parameters(), len(self.train_loader), {"epsilon": epsilon}

    def evaluate(self, parameters, config):
        self.set_parameters(parameters)
        self.model.eval()
        correct, total = 0, 0
        with torch.no_grad():
            for x, y in self.val_loader:
                correct += (self.model(x).argmax(1) == y).sum().item()
                total += len(y)
        return float(correct / total), total, {"accuracy": correct / total}