logoChat01.ai

#!/usr/bin/env python # -*- coding: utf-8 -*- """ ...

Erstellt am: 30. Oktober 2026

Beantwortet mit GPT-5.6 Thinking von Chat01

Frage

#!/usr/bin/env python

-- coding: utf-8 --

"""
Personalized Federated Anomaly Detection Benchmark
Replacing FedTreeLoRA (NLI) with Anomaly Detection Methods

This benchmark framework provides a comprehensive comparison of personalized
federated learning methods for anomaly detection tasks, structured similarly
to the FedTreeLoRA benchmark but adapted for anomaly detection.

Original FedTreeLoRA: https://github.com/jmbian/FedTreeLoRA
Our Adaptation: Anomaly detection with multi-domain heterogeneity

Baseline Methods Implemented:

  1. FedAvg - Standard federated averaging (McMahan et al., 2017)
  2. FedProx - Proximal regularization (Li et al., 2020)
  3. FedPer - Personalized layers (Arivazhagan et al., 2019)
  4. FedRep - Representation learning (Collins et al., 2021)
  5. Ditto - Personalized federated learning (Li et al., 2021)
  6. FedCR - Causal representation learning (our method)
  7. PFLAD - Random client selection baseline

Datasets:

  • Multi-domain anomaly detection (CalIt2, GECCO, Genesis, etc.)
  • Dirichlet distribution for statistical heterogeneity
  • Domain-specific assignments for functional heterogeneity

Evaluation Metrics:

  • ROC-AUC (Area Under ROC Curve)
  • PR-AUC (Precision-Recall AUC)
  • F1-Score
  • Per-domain performance
  • Communication efficiency

Usage:
# Run full benchmark
python federated_anomaly_benchmark.py --config benchmark_config.json

text
# Run specific method
python federated_anomaly_benchmark.py --method FedCR --num_clients 20

# Hyperparameter sensitivity
python federated_anomaly_benchmark.py --experiment hyperparameter_sweep

"""

import os
import json
import copy
import random
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Subset
from typing import Dict, List, Tuple, Optional, Any
from collections import defaultdict
import time
from tqdm import tqdm
from sklearn.metrics import roc_auc_score, average_precision_score, f1_score, precision_recall_curve
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

Import our existing modules

from federated_domain_manager import FederatedDomainManager
from dataset import *

============================================================================

Base Classes for Anomaly Detection Models

============================================================================

class AnomalyDetectorAutoencoder(nn.Module):
"""
Standard Autoencoder for anomaly detection.
Used as base model for all methods.
"""
def init(self, input_dim: int, hidden_dims: List[int], latent_dim: int):
super().init()
self.input_dim = input_dim
self.latent_dim = latent_dim

text
    # Encoder
    encoder_layers = []
    prev_dim = input_dim
    for h_dim in hidden_dims:
        encoder_layers.extend([
            nn.Linear(prev_dim, h_dim),
            nn.ReLU(),
            nn.Dropout(0.1)
        ])
        prev_dim = h_dim
    encoder_layers.append(nn.Linear(prev_dim, latent_dim))
    self.encoder = nn.Sequential(*encoder_layers)

    # Decoder
    decoder_layers = []
    prev_dim = latent_dim
    for h_dim in reversed(hidden_dims):
        decoder_layers.extend([
            nn.Linear(prev_dim, h_dim),
            nn.ReLU(),
            nn.Dropout(0.1)
        ])
        prev_dim = h_dim
    decoder_layers.append(nn.Linear(prev_dim, input_dim))
    self.decoder = nn.Sequential(*decoder_layers)

def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """Forward pass returning reconstruction and latent representation."""
    latent = self.encoder(x)
    reconstruction = self.decoder(latent)
    return reconstruction, latent

def compute_anomaly_score(self, x: torch.Tensor) -> torch.Tensor:
    """Compute reconstruction-based anomaly score."""
    reconstruction, _ = self.forward(x)
    mse = torch.mean((x - reconstruction) ** 2, dim=tuple(range(1, x.dim())))
    return mse

class PersonalizedAutoencoder(nn.Module):
"""
Autoencoder with separate global and personalized layers.
Used for FedPer and FedRep.
"""
def init(self, input_dim: int, hidden_dims: List[int], latent_dim: int,
personalized_layers: int = 2):
super().init()
self.input_dim = input_dim
self.latent_dim = latent_dim
self.personalized_layers = personalized_layers

text
    # Global encoder layers
    global_encoder = []
    prev_dim = input_dim
    for i, h_dim in enumerate(hidden_dims[:-personalized_layers]):
        global_encoder.extend([
            nn.Linear(prev_dim, h_dim),
            nn.ReLU(),
            nn.Dropout(0.1)
        ])
        prev_dim = h_dim
    self.global_encoder = nn.Sequential(*global_encoder)

    # Personalized encoder layers
    personal_encoder = []
    for h_dim in hidden_dims[-personalized_layers:]:
        personal_encoder.extend([
            nn.Linear(prev_dim, h_dim),
            nn.ReLU(),
            nn.Dropout(0.1)
        ])
        prev_dim = h_dim
    personal_encoder.append(nn.Linear(prev_dim, latent_dim))
    self.personal_encoder = nn.Sequential(*personal_encoder)

    # Decoder (can be global or personalized based on method)
    decoder_layers = []
    prev_dim = latent_dim
    for h_dim in reversed(hidden_dims):
        decoder_layers.extend([
            nn.Linear(prev_dim, h_dim),
            nn.ReLU(),
            nn.Dropout(0.1)
        ])
        prev_dim = h_dim
    decoder_layers.append(nn.Linear(prev_dim, input_dim))
    self.decoder = nn.Sequential(*decoder_layers)

def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    global_features = self.global_encoder(x)
    latent = self.personal_encoder(global_features)
    reconstruction = self.decoder(latent)
    return reconstruction, latent

def compute_anomaly_score(self, x: torch.Tensor) -> torch.Tensor:
    reconstruction, _ = self.forward(x)
    mse = torch.mean((x - reconstruction) ** 2, dim=tuple(range(1, x.dim())))
    return mse

============================================================================

Baseline Method Implementations

============================================================================

class BaselineClient:
"""Base client class for all personalized FL methods."""

text
def __init__(self, client_id: int, model: nn.Module, train_loader: DataLoader,
             val_loader: DataLoader, test_loader: DataLoader,
             device: str = 'cuda', lr: float = 1e-3, local_epochs: int = 5,
             domain_name: str = 'unknown'):
    self.client_id = client_id
    self.model = model.to(device)
    self.train_loader = train_loader
    self.val_loader = val_loader
    self.test_loader = test_loader
    self.device = device
    self.lr = lr
    self.local_epochs = local_epochs
    self.domain_name = domain_name
    self.optimizer = optim.Adam(self.model.parameters(), lr=lr)

def train_local(self, global_model_state: Optional[Dict] = None) -> Dict[str, torch.Tensor]:
    """Perform local training. Override in subclasses for specific methods."""
    self.model.train()

    for epoch in range(self.local_epochs):
        for batch_data, _ in self.train_loader:
            batch_data = batch_data.to(self.device)

            # Handle windowed data: [batch_size, win_size, features] -> [batch_size * win_size, features]
            if len(batch_data.shape) == 3:
                batch_size, win_size, features = batch_data.shape
                batch_data = batch_data.reshape(batch_size * win_size, features)
            elif len(batch_data.shape) == 2:
                # Already correct shape [batch_size, features]
                pass
            else:
                raise ValueError(f"Unexpected batch shape: {batch_data.shape}")

            self.optimizer.zero_grad()
            reconstruction, _ = self.model(batch_data)
            loss = nn.functional.mse_loss(reconstruction, batch_data)
            loss.backward()
            self.optimizer.step()

    return {k: v.cpu().clone() for k, v in self.model.state_dict().items()}

def evaluate(self, data_loader: Optional[DataLoader] = None) -> Dict[str, float]:
    """Evaluate anomaly detection performance."""
    self.model.eval()
    data_loader = data_loader or self.test_loader

    all_scores = []
    all_labels = []

    with torch.no_grad():
        for batch_data, batch_labels in data_loader:
            batch_data = batch_data.to(self.device)

            # Handle windowed data: [batch_size, win_size, features] -> [batch_size * win_size, features]
            original_batch_size = batch_data.shape[0]
            if len(batch_data.shape) == 3:
                batch_size, win_size, features = batch_data.shape
                batch_data = batch_data.reshape(batch_size * win_size, features)
                # Flatten labels similarly
                if len(batch_labels.shape) == 2:
                    batch_labels = batch_labels.reshape(-1)

            scores = self.model.compute_anomaly_score(batch_data)
            all_scores.extend(scores.cpu().numpy())
            all_labels.extend(batch_labels.numpy())

    all_scores = np.array(all_scores)
    all_labels = np.array(all_labels).flatten()

    # Compute metrics
    try:
        roc_auc = roc_auc_score(all_labels, all_scores)
        pr_auc = average_precision_score(all_labels, all_scores)

        # Compute F1 at best threshold
        precision, recall, thresholds = precision_recall_curve(all_labels, all_scores)
        f1_scores = 2 * (precision * recall) / (precision + recall + 1e-10)
        best_f1 = np.max(f1_scores)

        return {
            'roc_auc': roc_auc,
            'pr_auc': pr_auc,
            'f1_score': best_f1,
            'num_samples': len(all_labels)
        }
    except Exception as e:
        print(f"Warning: Could not compute metrics for client {self.client_id}: {e}")
        return {
            'roc_auc': 0.0,
            'pr_auc': 0.0,
            'f1_score': 0.0,
            'num_samples': len(all_labels)
        }

class FedAvgClient(BaselineClient):
"""FedAvg: Standard federated averaging."""

text
def train_local(self, global_model_state: Optional[Dict] = None) -> Dict[str, torch.Tensor]:
    # Synchronize with global model
    if global_model_state is not None:
        self.model.load_state_dict(global_model_state)

    # Standard local training
    return super().train_local(global_model_state)

class FedProxClient(BaselineClient):
"""FedProx: Federated learning with proximal term."""

text
def __init__(self, *args, mu: float = 0.01, **kwargs):
    super().__init__(*args, **kwargs)
    self.mu = mu
    self.global_model_state = None

def train_local(self, global_model_state: Optional[Dict] = None) -> Dict[str, torch.Tensor]:
    if global_model_state is not None:
        self.global_model_state = copy.deepcopy(global_model_state)
        self.model.load_state_dict(global_model_state)

    self.model.train()

    for epoch in range(self.local_epochs):
        for batch_data, _ in self.train_loader:
            batch_data = batch_data.to(self.device)

            self.optimizer.zero_grad()
            reconstruction, _ = self.model(batch_data)
            recon_loss = nn.functional.mse_loss(reconstruction, batch_data)

            # Proximal term
            if self.global_model_state is not None:
                prox_term = 0.0
                for name, param in self.model.named_parameters():
                    prox_term += torch.sum((param - self.global_model_state[name].to(self.device)) ** 2)
                loss = recon_loss + (self.mu / 2) * prox_term
            else:
                loss = recon_loss

            loss.backward()
            self.optimizer.step()

    return {k: v.cpu().clone() for k, v in self.model.state_dict().items()}

class FedPerClient(BaselineClient):
"""FedPer: Personalized layers remain local, only global layers are aggregated."""

text
def __init__(self, *args, personalized_layer_names: List[str], **kwargs):
    super().__init__(*args, **kwargs)
    self.personalized_layer_names = personalized_layer_names

def train_local(self, global_model_state: Optional[Dict] = None) -> Dict[str, torch.Tensor]:
    # Only update global layers from server
    if global_model_state is not None:
        current_state = self.model.state_dict()
        for name, param in global_model_state.items():
            if not any(pname in name for pname in self.personalized_layer_names):
                current_state[name] = param
        self.model.load_state_dict(current_state)

    # Train all parameters
    super().train_local(global_model_state)

    # Only return global layers for aggregation
    model_state = self.model.state_dict()
    global_state = {k: v.cpu().clone() for k, v in model_state.items()
                   if not any(pname in k for pname in self.personalized_layer_names)}

    return global_state

class FedRepClient(BaselineClient):
"""FedRep: Learn shared representations, personalize classifier."""

text
def __init__(self, *args, representation_layers: List[str], **kwargs):
    super().__init__(*args, **kwargs)
    self.representation_layers = representation_layers

def train_local(self, global_model_state: Optional[Dict] = None) -> Dict[str, torch.Tensor]:
    # Synchronize representation layers
    if global_model_state is not None:
        current_state = self.model.state_dict()
        for name, param in global_model_state.items():
            if any(rname in name for rname in self.representation_layers):
                current_state[name] = param
        self.model.load_state_dict(current_state)

    # Train all parameters
    super().train_local(global_model_state)

    # Only return representation layers for aggregation
    model_state = self.model.state_dict()
    rep_state = {k: v.cpu().clone() for k, v in model_state.items()
                if any(rname in k for rname in self.representation_layers)}

    return rep_state

class DittoClient(BaselineClient):
"""Ditto: Personalized federated learning with separate global and local models."""

text
def __init__(self, *args, lam: float = 0.1, **kwargs):
    super().__init__(*args, **kwargs)
    self.lam = lam
    self.global_model = copy.deepcopy(self.model)
    self.local_model = copy.deepcopy(self.model)

def train_local(self, global_model_state: Optional[Dict] = None) -> Dict[str, torch.Tensor]:
    # Update global model
    if global_model_state is not None:
        self.global_model.load_state_dict(global_model_state)

    # Train local model with regularization toward global model
    self.local_model.train()
    local_optimizer = optim.Adam(self.local_model.parameters(), lr=self.lr)

    for epoch in range(self.local_epochs):
        for batch_data, _ in self.train_loader:
            batch_data = batch_data.to(self.device)

            # Handle windowed data: [batch_size, win_size, features] -> [batch_size * win_size, features]
            if len(batch_data.shape) == 3:
                batch_size, win_size, features = batch_data.shape
                batch_data = batch_data.reshape(batch_size * win_size, features)
            elif len(batch_data.shape) == 2:
                # Already correct shape [batch_size, features]
                pass
            else:
                raise ValueError(f"Unexpected batch shape: {batch_data.shape}")

            local_optimizer.zero_grad()
            reconstruction, _ = self.local_model(batch_data)
            recon_loss = nn.functional.mse_loss(reconstruction, batch_data)

            # Regularization toward global model
            reg_term = 0.0
            for (n1, p1), (n2, p2) in zip(self.local_model.named_parameters(),
                                           self.global_model.named_parameters()):
                reg_term += torch.sum((p1 - p2) ** 2)

            loss = recon_loss + (self.lam / 2) * reg_term
            loss.backward()
            local_optimizer.step()

    # Use local model for evaluation
    self.model = self.local_model

    # Return global model state for aggregation
    return {k: v.cpu().clone() for k, v in self.global_model.state_dict().items()}

class FedCRClient(BaselineClient):
"""FedCR: Our causal representation learning method."""

text
def __init__(self, *args, alpha: float = 1.0, dimZ: int = 128, **kwargs):
    super().__init__(*args, **kwargs)
    self.alpha = alpha
    self.dimZ = dimZ

def train_local(self, global_model_state: Optional[Dict] = None) -> Dict[str, torch.Tensor]:
    if global_model_state is not None:
        self.model.load_state_dict(global_model_state)

    self.model.train()

    for epoch in range(self.local_epochs):
        for batch_data, _ in self.train_loader:
            batch_data = batch_data.to(self.device)

            self.optimizer.zero_grad()
            reconstruction, latent = self.model(batch_data)

            # Reconstruction loss
            recon_loss = nn.functional.mse_loss(reconstruction, batch_data)

            # Causal regularization (simplified for this implementation)
            # In full FedCR, this includes VIB and POE
            causal_reg = self.alpha * torch.mean(latent ** 2)

            loss = recon_loss + causal_reg
            loss.backward()
            self.optimizer.step()

    return {k: v.cpu().clone() for k, v in self.model.state_dict().items()}

============================================================================

Server Implementation

============================================================================

class FederatedServer:
"""Federated server for aggregation."""

text
def __init__(self, global_model: nn.Module, device: str = 'cuda'):
    self.global_model = global_model.to(device)
    self.device = device

def aggregate(self, client_updates: List[Dict[str, torch.Tensor]],
              client_weights: Optional[List[float]] = None) -> Dict[str, torch.Tensor]:
    """Federated averaging of client updates."""
    if not client_updates:
        return self.global_model.state_dict()

    if client_weights is None:
        client_weights = [1.0 / len(client_updates)] * len(client_updates)

    # Weighted average
    aggregated_state = {}
    for key in client_updates[0].keys():
        aggregated_state[key] = torch.sum(
            torch.stack([w * client_updates[i][key] for i, w in enumerate(client_weights)]),
            dim=0
        )

    self.global_model.load_state_dict(aggregated_state)
    return aggregated_state

============================================================================

Experiment Runner

============================================================================

class AnomalyDetectionBenchmark:
"""Main benchmark class for running experiments."""

text
def __init__(self, config: Dict[str, Any]):
    self.config = config
    self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    self.results = defaultdict(list)

    # Set random seeds
    self.set_seeds(config.get('seed', 42))

def set_seeds(self, seed: int):
    """Set random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)

def create_clients(self, method: str, num_clients: int,
                  domain_preset: str = 'diverse_6_clients') -> Tuple[List, FederatedServer]:
    """Create clients and server for a specific method."""

    # Load data using domain manager
    # If preset doesn't have enough clients, use auto_assign mode instead
    try:
        manager = FederatedDomainManager(preset=domain_preset)
        all_loaders = manager.get_all_client_loaders(batch_size=self.config['batch_size'])

        # If preset has fewer clients than requested, use auto_assign
        if len(all_loaders) < num_clients:
            print(f"Preset '{domain_preset}' only has {len(all_loaders)} clients, switching to auto_assign for {num_clients} clients")
            manager = FederatedDomainManager(num_clients=num_clients, auto_assign=True)
            all_loaders = manager.get_all_client_loaders(batch_size=self.config['batch_size'])
    except Exception as e:
        print(f"Error with preset '{domain_preset}': {e}")
        print(f"Falling back to auto_assign mode with {num_clients} clients")
        manager = FederatedDomainManager(num_clients=num_clients, auto_assign=True)
        all_loaders = manager.get_all_client_loaders(batch_size=self.config['batch_size'])

    # Adjust num_clients if it exceeds available clients
    actual_num_clients = min(num_clients, len(all_loaders))
    if actual_num_clients < num_clients:
        print(f"Warning: Requested {num_clients} clients but only {actual_num_clients} available. Using {actual_num_clients}.")

    # Detect input dimension from first client's data
    first_client_id = sorted(all_loaders.keys())[0]
    first_batch = next(iter(all_loaders[first_client_id]['train']))
    if isinstance(first_batch, (tuple, list)):
        sample_data = first_batch[0]
    else:
        sample_data = first_batch

    # Handle different data shapes
    if len(sample_data.shape) == 1:
        detected_input_dim = sample_data.shape[0]
    elif len(sample_data.shape) == 2:
        # [batch_size, features]
        detected_input_dim = sample_data.shape[1]
    elif len(sample_data.shape) == 3:
        # [batch_size, win_size, features] - we reshape to [batch_size * win_size, features] during training
        # So input_dim should be the last dimension (features), not win_size * features
        detected_input_dim = sample_data.shape[2]
    else:
        # Flatten to last dimension
        detected_input_dim = int(np.prod(sample_data.shape[1:]))

    print(f"Detected input dimension: {detected_input_dim} from shape {sample_data.shape}")

    # Override model config with detected dimension
    model_config = self.config['model'].copy()
    model_config['input_dim'] = detected_input_dim

    # Create base model
    base_model = AnomalyDetectorAutoencoder(
        input_dim=model_config['input_dim'],
        hidden_dims=model_config['hidden_dims'],
        latent_dim=model_config['latent_dim']
    )

    # Create clients based on method
    clients = []
    for client_id in sorted(all_loaders.keys())[:actual_num_clients]:
        loaders = all_loaders[client_id]
        metadata = loaders['metadata']

        client_model = copy.deepcopy(base_model)

        if method == 'FedAvg':
            client = FedAvgClient(
                client_id, client_model,
                loaders['train'], loaders['val'], loaders['test'],
                device=self.device,
                lr=self.config['lr'],
                local_epochs=self.config['local_epochs'],
                domain_name=metadata['domain_name']
            )
        elif method == 'FedProx':
            client = FedProxClient(
                client_id, client_model,
                loaders['train'], loaders['val'], loaders['test'],
                device=self.device,
                lr=self.config['lr'],
                local_epochs=self.config['local_epochs'],
                domain_name=metadata['domain_name'],
                mu=self.config.get('mu', 0.01)
            )
        elif method == 'FedPer':
            client = FedPerClient(
                client_id, client_model,
                loaders['train'], loaders['val'], loaders['test'],
                device=self.device,
                lr=self.config['lr'],
                local_epochs=self.config['local_epochs'],
                domain_name=metadata['domain_name'],
                personalized_layer_names=['personal_encoder', 'decoder']
            )
        elif method == 'FedRep':
            client = FedRepClient(
                client_id, client_model,
                loaders['train'], loaders['val'], loaders['test'],
                device=self.device,
                lr=self.config['lr'],
                local_epochs=self.config['local_epochs'],
                domain_name=metadata['domain_name'],
                representation_layers=['global_encoder']
            )
        elif method == 'Ditto':
            client = DittoClient(
                client_id, client_model,
                loaders['train'], loaders['val'], loaders['test'],
                device=self.device,
                lr=self.config['lr'],
                local_epochs=self.config['local_epochs'],
                domain_name=metadata['domain_name'],
                lam=self.config.get('lam', 0.1)
            )
        elif method == 'FedCR':
            client = FedCRClient(
                client_id, client_model,
                loaders['train'], loaders['val'], loaders['test'],
                device=self.device,
                lr=self.config['lr'],
                local_epochs=self.config['local_epochs'],
                domain_name=metadata['domain_name'],
                alpha=self.config.get('alpha', 1.0),
                dimZ=self.config.get('dimZ', 128)
            )
        else:
            raise ValueError(f"Unknown method: {method}")

        clients.append(client)

    # Create server
    server = FederatedServer(copy.deepcopy(base_model), device=self.device)

    return clients, server

def run_experiment(self, method: str, num_clients: int = 6,
                  global_rounds: int = 30, client_frac: float = 1.0) -> Dict:
    """Run a single experiment."""

    print(f"\n{'='*80}")
    print(f"Running {method} with {num_clients} clients for {global_rounds} rounds")
    print(f"{'='*80}\n")

    # Create clients and server
    clients, server = self.create_clients(method, num_clients)

    # Training loop
    round_results = []

    for round_idx in tqdm(range(global_rounds), desc=f"{method} Training"):
        # Client selection
        num_selected = max(1, int(client_frac * len(clients)))
        selected_clients = random.sample(clients, num_selected)

        # Local training
        client_updates = []
        client_weights = []

        for client in selected_clients:
            global_model_state = server.global_model.state_dict()
            update = client.train_local(global_model_state)
            client_updates.append(update)
            client_weights.append(1.0 / len(selected_clients))

        # Server aggregation
        server.aggregate(client_updates, client_weights)

        # Evaluation (every 5 rounds)
        if round_idx % 5 == 0 or round_idx == global_rounds - 1:
            eval_results = self.evaluate_clients(clients, round_idx)
            round_results.append(eval_results)

    return {
        'method': method,
        'num_clients': num_clients,
        'round_results': round_results,
        'final_results': round_results[-1] if round_results else None
    }

def evaluate_clients(self, clients: List, round_idx: int) -> Dict:
    """Evaluate all clients."""
    all_metrics = []
    domain_metrics = defaultdict(list)

    for client in clients:
        metrics = client.evaluate()
        metrics['client_id'] = client.client_id
        metrics['domain'] = client.domain_name
        metrics['round'] = round_idx

        all_metrics.append(metrics)
        domain_metrics[client.domain_name].append(metrics)

    # Aggregate results
    avg_metrics = {
        'round': round_idx,
        'roc_auc_mean': np.mean([m['roc_auc'] for m in all_metrics]),
        'roc_auc_std': np.std([m['roc_auc'] for m in all_metrics]),
        'pr_auc_mean': np.mean([m['pr_auc'] for m in all_metrics]),
        'pr_auc_std': np.std([m['pr_auc'] for m in all_metrics]),
        'f1_mean': np.mean([m['f1_score'] for m in all_metrics]),
        'f1_std': np.std([m['f1_score'] for m in all_metrics]),
        'per_client': all_metrics,
        'per_domain': dict(domain_metrics)
    }

    print(f"\nRound {round_idx}: ROC-AUC={avg_metrics['roc_auc_mean']:.4f}±{avg_metrics['roc_auc_std']:.4f}, "
          f"F1={avg_metrics['f1_mean']:.4f}±{avg_metrics['f1_std']:.4f}")

    return avg_metrics

def run_full_benchmark(self, methods: List[str], save_dir: str = './results'):
    """Run full benchmark comparing all methods."""

    os.makedirs(save_dir, exist_ok=True)

    all_results = {}

    for method in methods:
        try:
            result = self.run_experiment(
                method=method,
                num_clients=self.config['num_clients'],
                global_rounds=self.config['global_rounds'],
                client_frac=self.config.get('client_frac', 1.0)
            )
            all_results[method] = result

            # Save individual result
            with open(os.path.join(save_dir, f'{method}_results.json'), 'w') as f:
                json.dump(result, f, indent=2, default=str)

        except Exception as e:
            print(f"\nError running {method}: {e}")
            import traceback
            traceback.print_exc()

    # Save all results
    with open(os.path.join(save_dir, 'all_results.json'), 'w') as f:
        json.dump(all_results, f, indent=2, default=str)

    # Generate comparison plots
    self.plot_comparison(all_results, save_dir)

    return all_results

def plot_comparison(self, results: Dict, save_dir: str):
    """Generate comparison plots."""

    # Extract data for plotting
    methods = list(results.keys())
    final_metrics = {}

    for method, result in results.items():
        if result['final_results']:
            final_metrics[method] = result['final_results']

    if not final_metrics:
        print("No results to plot")
        return

    # Plot ROC-AUC comparison
    plt.figure(figsize=(12, 6))

    # Bar plot
    plt.subplot(1, 2, 1)
    method_names = list(final_metrics.keys())
    roc_aucs = [final_metrics[m]['roc_auc_mean'] for m in method_names]
    roc_stds = [final_metrics[m]['roc_auc_std'] for m in method_names]

    plt.bar(method_names, roc_aucs, yerr=roc_stds, capsize=5)
    plt.ylabel('ROC-AUC')
    plt.title('Final ROC-AUC Comparison')
    plt.xticks(rotation=45)
    plt.grid(axis='y', alpha=0.3)

    # Learning curves
    plt.subplot(1, 2, 2)
    for method, result in results.items():
        rounds = [r['round'] for r in result['round_results']]
        roc_aucs = [r['roc_auc_mean'] for r in result['round_results']]
        plt.plot(rounds, roc_aucs, marker='o', label=method)

    plt.xlabel('Communication Round')
    plt.ylabel('ROC-AUC')
    plt.title('Learning Curves')
    plt.legend()
    plt.grid(alpha=0.3)

    plt.tight_layout()
    plt.savefig(os.path.join(save_dir, 'comparison_plots.png'), dpi=300, bbox_inches='tight')
    plt.close()

    print(f"\nPlots saved to {save_dir}/comparison_plots.png")

============================================================================

Main Entry Point

============================================================================

def parse_args():
parser = argparse.ArgumentParser(description='Personalized Federated Anomaly Detection Benchmark')

text
parser.add_argument('--config', type=str, default=None,
                   help='Path to config JSON file')
parser.add_argument('--method', type=str, default='all',
                   choices=['all', 'FedAvg', 'FedProx', 'FedPer', 'FedRep', 'Ditto', 'FedCR'],
                   help='Method to run')
parser.add_argument('--num_clients', type=int, default=6,
                   help='Number of clients')
parser.add_argument('--global_rounds', type=int, default=30,
                   help='Number of global communication rounds')
parser.add_argument('--local_epochs', type=int, default=5,
                   help='Number of local epochs')
parser.add_argument('--batch_size', type=int, default=32,
                   help='Batch size')
parser.add_argument('--lr', type=float, default=1e-3,
                   help='Learning rate')
parser.add_argument('--seed', type=int, default=42,
                   help='Random seed')
parser.add_argument('--save_dir', type=str, default='./benchmark_results',
                   help='Directory to save results')

return parser.parse_args()

def main():
args = parse_args()

text
# Load or create config
if args.config:
    with open(args.config, 'r') as f:
        config = json.load(f)
else:
    config = {
        'seed': args.seed,
        'num_clients': args.num_clients,
        'global_rounds': args.global_rounds,
        'local_epochs': args.local_epochs,
        'batch_size': args.batch_size,
        'lr': args.lr,
        'model': {
            'input_dim': 64 * 64,  # Adjust based on your data
            'hidden_dims': [256, 128],
            'latent_dim': 64
        },
        'mu': 0.01,  # FedProx
        'lam': 0.1,  # Ditto
        'alpha': 1.0,  # FedCR
        'dimZ': 128  # FedCR
    }

# Create benchmark
benchmark = AnomalyDetectionBenchmark(config)

# Run experiments
if args.method == 'all':
    methods = ['FedAvg', 'FedProx', 'FedPer', 'FedRep', 'Ditto', 'FedCR']
    benchmark.run_full_benchmark(methods, args.save_dir)
else:
    result = benchmark.run_experiment(
        method=args.method,
        num_clients=args.num_clients,
        global_rounds=args.global_rounds
    )
    os.makedirs(args.save_dir, exist_ok=True)
    with open(os.path.join(args.save_dir, f'{args.method}_results.json'), 'w') as f:
        json.dump(result, f, indent=2, default=str)

print(f"\n{'='*80}")
print("Benchmark Complete!")
print(f"Results saved to {args.save_dir}")
print(f"{'='*80}\n")

if name == 'main':
main() 为什么一直报错呀Detected input dimension: 1 from shape torch.Size([32, 32, 1])
FedCR Training: 0%| | 0/30 [00:00<?, ?it/s]

Error running FedCR: mat1 and mat2 shapes cannot be multiplied (2048x38 and 1x256)
Traceback (most recent call last):
File "/tmp/pycharm_project_757/code_FedCR/federated_anomaly_benchmark.py", line 757, in run_full_benchmark
result = self.run_experiment(
File "/tmp/pycharm_project_757/code_FedCR/federated_anomaly_benchmark.py", line 697, in run_experiment
update = client.train_local(global_model_state)
File "/tmp/pycharm_project_757/code_FedCR/federated_anomaly_benchmark.py", line 462, in train_local
reconstruction, latent = self.model(batch_data)
File "/home/zmn/miniconda3/envs/cpiri/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/zmn/miniconda3/envs/cpiri/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
return forward_call(*args, **kwargs)
File "/tmp/pycharm_project_757/code_FedCR/federated_anomaly_benchmark.py", line 112, in forward
latent = self.encoder(x)
File "/home/zmn/miniconda3/envs/cpiri/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/zmn/miniconda3/envs/cpiri/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
return forward_call(*args, **kwargs)
File "/home/zmn/miniconda3/envs/cpiri/lib/python3.10/site-packages/torch/nn/modules/container.py", line 250, in forward
input = module(input)
File "/home/zmn/miniconda3/envs/cpiri/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/zmn/miniconda3/envs/cpiri/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
return forward_call(*args, **kwargs)
File "/home/zmn/miniconda3/envs/cpiri/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 134, in forward
return F.linear(input, self.weight, self.bias)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (2048x38 and 1x256)
No results to plot

Gedanke

Diagnosing dimension mismatch

I’m tracing the error to inconsistent feature widths across clients: the model is initialized with input size 1, but another batch has 38 features. The shared autoencoder therefore cannot process all client data without standardization.

Diagnosed dimension mismatch

Antwort

这个报错的核心不是 PyTorch 本身，而是不同客户端的数据特征维度不一致，但你只用第一个客户端的维度创建了全局模型。另外 FedCRClient.train_local() 里也漏掉了你在基类中已经写好的 3D→2D 处理。

Gedanke

Spotting further issues

I’m also seeing architectural mismatches: the representation-layer names differ, the aggregation may collect nothing, and the personalized training returns an unchanged global model. These require separate fixes after resolving input dimensions.

Spotted further issues

Antwort

从报错数字就能反推出问题：

  • 你先打印的是 torch.Size([32, 32, 1])，所以代码把 input_dim 设成了 1，第一层变成 Linear(1, 256)。
  • 真正报错时 mat1 是 2048 x 38。2048 = 32 × 64，这很像某个客户端实际 batch 是 [32, 64, 38]。
  • nn.Linear 会把最后一维当成特征维，因此它实际收到 38 个特征，却遇到只接受 1 个特征的 Linear(1,256)，于是报： mat1 and mat2 shapes cannot be multiplied (...x38 and 1x256)。

所以即使你把 FedCR 的 3D 数据 reshape 成二维，最终也仍然是 [..., 38]，单纯 reshape 不能解决 1 vs 38 的冲突。

Gedanke

Worked for 39s

Antwort

问题已经很明确了：你的不同 domain/client 的原始特征维度不同，但全局 Autoencoder 只根据第一个 client 的维度初始化。

你的这两行日志已经把问题暴露出来了：

text
Detected input dimension: 1 from shape torch.Size([32, 32, 1])

因此模型第一层实际上是：

python
nn.Linear(1, 256)

但真正训练 FedCR 时，某个客户端传进模型的数据最后一维是 38：

text
mat1 and mat2 shapes cannot be multiplied (2048x38 and 1x256)

这里 2048 很可能就是：

text
batch_size × win_size
= 32 × 64
= 2048

也就是说这个 client 的数据大概率是：

python
[32, 64, 38]
# batch=32
# window=64
# features=38

PyTorch 的 Linear 永远把最后一维当 feature，所以实际发生的是：

python
[32, 64, 38]
       ↓ Linear
[..., 38] @ [1, 256]   # ❌ 38 != 1

而你的模型是按照另一个 client：

python
[32, 32, 1]

创建的。

你代码里真正的问题在这里

create_clients() 只检查了第一个客户端：

python
first_client_id = sorted(all_loaders.keys())[0]
first_batch = next(iter(all_loaders[first_client_id]['train']))

然后：

python
detected_input_dim = sample_data.shape[2]

所以只要 Client 0 是：

python
[B, W, 1]

你就直接决定：

python
input_dim = 1

但是后面的 Client 可能是：

python
Client 0: [B, 32, 1]
Client 1: [B, 64, 38]
Client 2: [B, 32, 18]
Client 3: [B, 100, 1]
...

这在你这种：

CalIt2, GECCO, Genesis 等不同 anomaly detection domain

里其实非常可能。

不同数据集本来就不一定具有相同数量的传感器/变量。

第一件事：先把所有客户端 shape 打出来

不要只检测第一个 client。

把你现在这段：

python
# Detect input dimension from first client's data
first_client_id = sorted(all_loaders.keys())[0]
first_batch = next(iter(all_loaders[first_client_id]['train']))
...

替换为下面这个诊断版本：

python
# ============================================================
# Inspect input dimensions of ALL clients
# ============================================================

client_feature_dims = {}

print("\n" + "=" * 80)
print("Client data shapes")
print("=" * 80)

for client_id in sorted(all_loaders.keys()):
    loaders = all_loaders[client_id]

    batch = next(iter(loaders['train']))

    if isinstance(batch, (tuple, list)):
        sample_data = batch[0]
    else:
        sample_data = batch

    if sample_data.ndim == 1:
        feature_dim = sample_data.shape[0]

    elif sample_data.ndim == 2:
        # [B, F]
        feature_dim = sample_data.shape[-1]

    elif sample_data.ndim == 3:
        # [B, W, F]
        feature_dim = sample_data.shape[-1]

    else:
        raise ValueError(
            f"Unsupported input shape for client {client_id}: "
            f"{sample_data.shape}"
        )

    domain_name = loaders.get('metadata', {}).get(
        'domain_name', 'unknown'
    )

    client_feature_dims[client_id] = feature_dim

    print(
        f"Client {client_id:2d} | "
        f"domain={domain_name:15s} | "
        f"shape={tuple(sample_data.shape)} | "
        f"feature_dim={feature_dim}"
    )

print("=" * 80)

unique_dims = sorted(set(client_feature_dims.values()))

print(f"Unique feature dimensions: {unique_dims}")

if len(unique_dims) != 1:
    raise ValueError(
        "\nDifferent clients have different feature dimensions!\n"
        f"{client_feature_dims}\n\n"
        "A single FedAvg-style autoencoder cannot aggregate models "
        "with different input/output dimensions."
    )

detected_input_dim = unique_dims[0]

你运行之后，我估计会看到类似：

text
Client  0 | domain=CalIt2    | shape=(32, 32, 1)  | feature_dim=1
Client  1 | domain=GECCO     | shape=(32, 64, 38) | feature_dim=38
Client  2 | domain=Genesis   | shape=(32, 32, 18) | feature_dim=18
...

Unique feature dimensions: [1, 18, 38]

这样就能彻底确认。

第二个 bug：FedCR 忘记处理 3D 数据

你的 BaselineClient.train_local() 已经写了：

python
if len(batch_data.shape) == 3:
    batch_size, win_size, features = batch_data.shape
    batch_data = batch_data.reshape(
        batch_size * win_size,
        features
    )

但是 FedCRClient.train_local() 又自己重写了一遍训练循环：

python
for batch_data, _ in self.train_loader:
    batch_data = batch_data.to(self.device)

    self.optimizer.zero_grad()
    reconstruction, latent = self.model(batch_data)

这里没有 reshape。

FedProxClient 也有同样的问题。

虽然这不是当前 1 vs 38 的根本原因，因为：

python
[32,64,38] -> [2048,38]

最后还是 38 维，但最好统一掉，否则不同算法的数据处理方式都不一样，benchmark 不公平。

我建议直接在 BaselineClient 写一个统一函数。

python
class BaselineClient:

    ...

    def prepare_batch(self, batch_data):
        batch_data = batch_data.to(self.device).float()

        if batch_data.ndim == 3:
            # [B, W, F] -> [B*W, F]
            batch_data = batch_data.reshape(
                -1,
                batch_data.shape[-1]
            )

        elif batch_data.ndim == 2:
            # [B, F]
            pass

        elif batch_data.ndim == 1:
            batch_data = batch_data.unsqueeze(0)

        else:
            raise ValueError(
                f"Client {self.client_id}, "
                f"domain={self.domain_name}: "
                f"unexpected input shape {batch_data.shape}"
            )

        expected_dim = self.model.input_dim
        actual_dim = batch_data.shape[-1]

        if actual_dim != expected_dim:
            raise ValueError(
                f"\nInput dimension mismatch!\n"
                f"Client: {self.client_id}\n"
                f"Domain: {self.domain_name}\n"
                f"Model expects: {expected_dim}\n"
                f"Data provides: {actual_dim}\n"
                f"Batch shape: {tuple(batch_data.shape)}"
            )

        return batch_data

然后所有算法都：

python
for batch_data, _ in self.train_loader:

    batch_data = self.prepare_batch(batch_data)

    self.optimizer.zero_grad()

    reconstruction, latent = self.model(batch_data)

    ...

包括：

python
BaselineClient
FedProxClient
FedCRClient
DittoClient

全部调用同一个函数。

这样以后错误会直接变成：

text
Input dimension mismatch!
Client: 3
Domain: GECCO
Model expects: 1
Data provides: 38
Batch shape: (2048, 38)

会比现在好调试很多。

但是，这仍然没有解决最根本的问题

你的 benchmark 声称研究：

multi-domain heterogeneity

可是现在不同 domain 连输入空间都不一样。

例如：

text
CalIt2       → 1 feature
GECCO        → 38 features
Genesis      → 18 features

那么这个模型：

python
Autoencoder(
    Linear(input_dim, 256),
    ...
    Linear(256, input_dim)
)

不能直接进行标准 FedAvg。

因为 Client A 需要：

python
Linear(1, 256)

Client B 需要：

python
Linear(38, 256)

参数 shape 分别是：

text
[256, 1]
[256, 38]

服务器根本没办法：

python
0.5 * weight_A + 0.5 * weight_B

它们 shape 都不同。

所以不能简单地：

python
每个 client 自动检测 input_dim

然后各自建立模型。

这样下一步 server.aggregate() 一样会挂。

所以你需要决定 benchmark 的建模方式

对于你目前这个项目，我更推荐下面这个方向：

text
Raw domain data
      ↓
Domain-specific preprocessing
      ↓
Same-dimensional representation
      ↓
Federated Autoencoder
      ↓
Anomaly score

也就是说让：

text
CalIt2  ─┐
GECCO   ─┼──> fixed dimension，例如 32/64
Genesis ─┘

之后所有 client 才能共享：

python
Linear(64, 256)

这样 FedAvg、FedProx、Ditto 才是数学上成立的。

暂时只想让代码先跑起来

最简单的办法是：

只选择 feature dimension 一样的 clients。

例如：

python
from collections import Counter

feature_dims = {}

for client_id, loaders in all_loaders.items():
    batch = next(iter(loaders['train']))
    x = batch[0] if isinstance(batch, (tuple, list)) else batch

    feature_dims[client_id] = x.shape[-1]

# 找出现最多的 feature dimension
counter = Counter(feature_dims.values())
target_dim = counter.most_common(1)[0][0]

print(f"Using feature dimension: {target_dim}")

eligible_clients = [
    client_id
    for client_id, dim in feature_dims.items()
    if dim == target_dim
]

print("Eligible clients:", eligible_clients)

all_loaders = {
    cid: all_loaders[cid]
    for cid in eligible_clients
}

例如最后只留下：

text
Client 0: feature=1
Client 3: feature=1
Client 5: feature=1

这样：

python
input_dim = 1

大家就一致了。

这是很好的 debug 方案，但不是你论文里最终 multi-domain benchmark 的理想方案。

不建议你直接 padding 到 38

你可能想到：

python
1 feature  -> padding -> 38
18 feature -> padding -> 38
38 feature -> 38

技术上能跑：

python
F.pad(x, (0, 38 - x.shape[-1]))

但不建议直接拿这个做最终实验。

因为：

text
GECCO feature 1

和：

text
CalIt2 feature 1

未必具有任何相同语义。

而且 reconstruction MSE 会把大量 padding 的 0 也统计进去：

python
MSE(
    predicted_38_dim,
    padded_38_dim
)

最后可能让 ROC-AUC/F1 出现很奇怪的结果。

如果只是验证代码 pipeline，可以 padding；正式 benchmark 最好不要。

你的代码还有几个隐藏 bug

你把当前这个 shape bug 修掉后，很可能马上又遇到别的问题。

FedRep 目前基本上一定有问题

你创建的是：

python
base_model = AnomalyDetectorAutoencoder(...)

它的 state dict 名字类似：

text
encoder.0.weight
encoder.0.bias
encoder.3.weight
decoder.0.weight
...

但是 FedRep 传的是：

python
representation_layers=['global_encoder']

然后：

python
if any(rname in k for rname in self.representation_layers)

实际上根本没有：

text
global_encoder.xxx

所以：

python
rep_state = {}

极有可能是空字典。

你其实定义了：

python
PersonalizedAutoencoder

但 create_clients() 完全没用它。

FedPer 也没有按你想的模型结构工作

你传的是：

python
personalized_layer_names=[
    'personal_encoder',
    'decoder'
]

但 AnomalyDetectorAutoencoder 里没有：

text
personal_encoder

只有：

text
encoder
decoder

所以实际上你只 personalized 了：

text
decoder

而不是代码注释里描述的结构。

FedPer/FedRep 的 server aggregation 还会有问题

你这里：

python
self.global_model.load_state_dict(aggregated_state)

默认：

python
strict=True

但 FedPer 返回的是 partial state：

python
{
    encoder....
}

没有 decoder。

所以很可能继续报：

text
Missing key(s) in state_dict

partial aggregation 应该是：

python
current_state = self.global_model.state_dict()

for key, value in aggregated_state.items():
    current_state[key] = value

self.global_model.load_state_dict(current_state)

而不是直接：

python
self.global_model.load_state_dict(aggregated_state)

Ditto 现在实际上没有训练 global model

这部分也很关键：

python
if global_model_state is not None:
    self.global_model.load_state_dict(global_model_state)

然后你只训练：

python
self.local_model

最后却返回：

python
return {
    k: v.cpu().clone()
    for k, v in self.global_model.state_dict().items()
}

也就是说：

text
server global model
        ↓
复制给 client
        ↓
global model 根本没更新
        ↓
原封不动传回 server

服务器平均了一堆没有变化的参数。

所以你的 Ditto global model 基本不会学习。

FedCR 的 dimZ 目前没有任何作用

这里：

python
self.dimZ = dimZ

之后就再也没有使用。

而且目前所谓：

python
causal_reg = self.alpha * torch.mean(latent ** 2)

本质上只是 latent L2 regularization：

text
L = reconstruction loss + α ||z||²

严格来说这还不能称为完整的 causal representation learning。

你注释里也写了：

python
# simplified
# full FedCR includes VIB and POE

如果后面要拿它做论文 baseline comparison，这部分最好特别注意。

我建议你现在先这么改

当前最优先的顺序是：

  1. 打印每个 client 的 domain + batch shape + feature_dim。
  2. 确认是不是出现了 1 / 18 / 38 等不同输入维度。
  3. 给所有 Client 加统一的 prepare_batch()。
  4. 调试阶段先只选相同 feature dimension 的 domains，把 FedCR 跑通。
  5. 再决定正式 benchmark 是做统一 feature representation，还是设计 domain-specific adapter。
  6. 之后修 FedPer / FedRep / Ditto 的实现。

所以这一次报错最核心的一句话就是：

text
你根据 Client 0 创建了 Linear(1,256)，
但 FedCR 后来抽到了一个每个时间点有 38 个特征的 Client，
于是 Linear 收到 38 维输入，直接 shape mismatch。

而不是简单的：

python
reshape 写错了

实际上 2048x38 已经非常明确地告诉你：真正冲突的是 38 != 1。

Diesen Q&A teilen