實作:攻擊聯邦學習

from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
import numpy as np
 
def load_partitioned_data(num_clients: int, data_dir: str = "./data"):
    """Partition MNIST into non-overlapping subsets for each client."""
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
 
    train_dataset = datasets.MNIST(data_dir, train=True, download=True,
                                    transform=transform)
    test_dataset = datasets.MNIST(data_dir, train=False, transform=transform)
 
    indices = np.random.permutation(len(train_dataset))
    partition_size = len(indices) // num_clients
    partitions = []
 
    for i in range(num_clients):
        start = i * partition_size
        end = start + partition_size
        client_indices = indices[start:end]
        client_dataset = Subset(train_dataset, client_indices)
        partitions.append(DataLoader(client_dataset, batch_size=32, shuffle=True))
 
    test_loader = DataLoader(test_dataset, batch_size=256)
    return partitions, test_loader

import flwr as fl
from flwr.server.strategy import FedAvg
 
def run_experiment(num_clients=10, num_malicious=2, num_rounds=20,
                   aggregation="fedavg", attack_type="label_flip"):
    """Run a federated learning experiment with attacks and defenses."""
    partitions, test_loader = load_partitioned_data(num_clients)
 
    def client_fn(cid: str):
        idx = int(cid)
        model = SimpleNet()
        if idx < num_malicious:
            return ByzantineClient(model, partitions[idx], idx, attack_type).to_client()
        else:
            return HonestClient(model, partitions[idx], idx).to_client()
 
    if aggregation == "fedavg":
        strategy = FedAvg(
            min_fit_clients=num_clients,
            min_evaluate_clients=num_clients,
            min_available_clients=num_clients,
        )
    elif aggregation == "trimmed_mean":
        strategy = TrimmedMeanStrategy(
            trim_ratio=0.2,
            min_fit_clients=num_clients,
        )
 
    history = fl.simulation.start_simulation(
        client_fn=client_fn,
        num_clients=num_clients,
        config=fl.server.ServerConfig(num_rounds=num_rounds),
        strategy=strategy,
    )
    return history
 
results = {}
for agg in ["fedavg", "trimmed_mean"]:
    for n_mal in [0, 1, 2, 3]:
        key = f"{agg}_mal{n_mal}"
        print(f"\n{'='*50}")
        print(f"Running: {key}")
        results[key] = run_experiment(
            num_clients=10, num_malicious=n_mal,
            aggregation=agg, attack_type="label_flip"
        )