怎样用Python实现基于联邦学习的跨设备异常检测？

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import flwr as fl
import collections

# 定义自编码器模型
class Autoencoder(nn.Module):
    def __init__(self, input_dim):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 16) # 编码到低维潜空间
        )
        self.decoder = nn.Sequential(
            nn.Linear(16, 32),
            nn.ReLU(),
            nn.Linear(32, 64),
            nn.ReLU(),
            nn.Linear(64, input_dim) # 解码回原始维度
        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded

# 辅助函数：获取模型参数
def get_parameters(net):
    return [val.cpu().numpy() for _, val in net.state_dict().items()]

# 辅助函数：设置模型参数
def set_parameters(net, parameters):
    params_dict = zip(net.state_dict().keys(), parameters)
    state_dict = collections.OrderedDict({k: torch.tensor(v) for k, v in params_dict})
    net.load_state_dict(state_dict, strict=True)

# 模拟数据生成（简化版，实际应用中数据来自设备）
def generate_client_data(num_samples=1000, input_dim=10, num_clients=3):
    all_data = np.random.rand(num_samples, input_dim).astype(np.float32)
    # 随机添加一些“异常”数据（例如，某些维度值特别大或小）
    num_anomalies = int(num_samples * 0.05)
    anomaly_indices = np.random.choice(num_samples, num_anomalies, replace=False)
    all_data[anomaly_indices] += np.random.rand(num_anomalies, input_dim) * 5 # 增加噪音

    # 将数据分割给不同的客户端
    client_data_list = np.array_split(all_data, num_clients)
    datasets = []
    for client_data in client_data_list:
        datasets.append(TensorDataset(torch.from_numpy(client_data)))
    return datasets

class AnomalyDetectionClient(fl.client.NumPyClient):
    def __init__(self, cid, net, trainloader):
        self.cid = cid
        self.net = net
        self.trainloader = trainloader
        self.criterion = nn.MSELoss() # 自编码器常用MSE作为重建误差
        self.optimizer = optim.Adam(self.net.parameters(), lr=0.001)

    def get_parameters(self, config):
        print(f"[Client {self.cid}] get_parameters")
        return get_parameters(self.net)

    def fit(self, parameters, config):
        print(f"[Client {self.cid}] fit, epoch: {config['local_epochs']}")
        set_parameters(self.net, parameters)
        # 局部训练
        self.net.train()
        for epoch in range(config["local_epochs"]):
            for batch_idx, (data,) in enumerate(self.trainloader):
                self.optimizer.zero_grad()
                outputs = self.net(data)
                loss = self.criterion(outputs, data)
                loss.backward()
                self.optimizer.step()
        print(f"[Client {self.cid}] local loss: {loss.item()}")
        return get_parameters(self.net), len(self.trainloader.dataset), {}

    def evaluate(self, parameters, config):
        print(f"[Client {self.cid}] evaluate")
        set_parameters(self.net, parameters)
        self.net.eval()
        total_loss = 0.0
        with torch.no_grad():
            for data, in self.trainloader: # 这里用训练集评估，实际可以有单独的测试集
                outputs = self.net(data)
                loss = self.criterion(outputs, data)
                total_loss += loss.item() * data.size(0)
        avg_loss = total_loss / len(self.trainloader.dataset)
        return avg_loss, len(self.trainloader.dataset), {"average_loss": avg_loss}

# 定义服务器端聚合策略
# 这里使用FedAvg策略，可以根据需求选择其他策略
strategy = fl.server.strategy.FedAvg(
    fraction_fit=1.0,  # 每次训练选择所有客户端
    fraction_evaluate=1.0, # 每次评估选择所有客户端
    min_fit_clients=2, # 至少需要2个客户端参与训练
    min_evaluate_clients=2, # 至少需要2个客户端参与评估
    min_available_clients=2, # 至少需要2个客户端在线
    evaluate_fn=None, # 服务器端不进行评估，由客户端完成
    on_fit_config_fn=lambda server_round: {"local_epochs": 5}, # 每轮训练客户端本地训练5个epoch
)

# 启动服务器
def start_server(num_rounds=5):
    fl.server.start_server(
        server_address="0.0.0.0:8080",
        config=fl.server.ServerConfig(num_rounds=num_rounds),
        strategy=strategy,
    )

# 模拟客户端启动（在实际中，这些会在不同的设备上运行）
def start_client(cid, data_set, input_dim):
    net = Autoencoder(input_dim)
    trainloader = DataLoader(data_set, batch_size=32, shuffle=True)
    client = AnomalyDetectionClient(cid, net, trainloader)
    fl.client.start_client(server_address="127.0.0.1:8080", client=client)

# 主程序入口
if __name__ == "__main__":
    INPUT_DIM = 10
    NUM_CLIENTS = 3
    NUM_ROUNDS = 5 # 联邦学习的轮次

    # 模拟生成数据
    client_datasets = generate_client_data(num_samples=1000, input_dim=INPUT_DIM, num_clients=NUM_CLIENTS)

    # 在单独的线程或进程中启动服务器和客户端
    # 这里为了演示方便，在同一个脚本中启动，实际部署需要分开
    import threading
    server_thread = threading.Thread(target=start_server, args=(NUM_ROUNDS,))
    server_thread.start()

    # 等待服务器启动
    import time
    time.sleep(5)

    client_threads = []
    for i in range(NUM_CLIENTS):
        client_thread = threading.Thread(target=start_client, args=(i, client_datasets[i], INPUT_DIM))
        client_threads.append(client_thread)
        client_thread.start()

    for t in client_threads:
        t.join()
    server_thread.join()
    print("联邦学习异常检测训练完成。")

    # 训练完成后，可以获取全局模型参数，并用于新的数据推理
    # 比如从服务器端保存模型，或让客户端加载最终模型进行推理
    # 这里省略了推理部分的实现，但核心是：
    # 1. 加载训练好的全局模型参数到Autoencoder实例
    # 2. 对新数据进行前向传播，计算重建误差
    # 3. 设置一个重建误差阈值，超过阈值的即为异常