|
|
"""aggregator.py module.""" |
|
|
|
|
|
import tensorflow as tf |
|
|
from typing import List, Dict |
|
|
import numpy as np |
|
|
from collections import defaultdict |
|
|
import logging |
|
|
|
|
|
class FederatedAggregator: |
|
|
def __init__(self, config: Dict): |
|
|
logger = logging.getLogger(__name__) |
|
|
logger.debug(f"Initializing FederatedAggregator with config: {config}") |
|
|
|
|
|
agg_config = None |
|
|
if 'aggregation' in config: |
|
|
agg_config = config['aggregation'] |
|
|
elif 'server' in config and 'aggregation' in config['server']: |
|
|
agg_config = config['server']['aggregation'] |
|
|
else: |
|
|
logger.error(f"No 'aggregation' key found in config passed to FederatedAggregator: {config}") |
|
|
raise KeyError("'aggregation' config section is required for FederatedAggregator") |
|
|
self.weighted = agg_config.get('weighted', True) |
|
|
logger.info(f"FederatedAggregator initialized. Weighted: {self.weighted}") |
|
|
|
|
|
def federated_averaging(self, updates: List[Dict]) -> List: |
|
|
"""Perform federated averaging (FedAvg) on model weights.""" |
|
|
logger = logging.getLogger(__name__) |
|
|
logger.info(f"Performing federated averaging on {len(updates)} client updates") |
|
|
|
|
|
if not updates: |
|
|
logger.warning("No updates provided for federated averaging") |
|
|
return None |
|
|
|
|
|
|
|
|
total_samples = sum(update['size'] for update in updates) |
|
|
logger.debug(f"Total samples across clients: {total_samples}") |
|
|
|
|
|
|
|
|
first_weights = updates[0]['weights'] |
|
|
aggregated_weights = [np.zeros_like(w) for w in first_weights] |
|
|
|
|
|
|
|
|
for update in updates: |
|
|
client_weights = update['weights'] |
|
|
client_size = update['size'] |
|
|
weight_factor = client_size / total_samples if self.weighted else 1.0 / len(updates) |
|
|
|
|
|
logger.debug(f"Client {update['client_id']}: size={client_size}, weight_factor={weight_factor}") |
|
|
|
|
|
|
|
|
for i, (agg_w, client_w) in enumerate(zip(aggregated_weights, client_weights)): |
|
|
aggregated_weights[i] += np.array(client_w) * weight_factor |
|
|
|
|
|
logger.info("Federated averaging completed successfully") |
|
|
return aggregated_weights |
|
|
|
|
|
def compute_metrics(self, client_metrics: List[Dict]) -> Dict: |
|
|
logger = logging.getLogger(__name__) |
|
|
logger.debug(f"Computing metrics for {len(client_metrics)} clients") |
|
|
if not client_metrics: |
|
|
logger.warning("No client metrics provided to compute_metrics.") |
|
|
return {} |
|
|
aggregated_metrics = defaultdict(float) |
|
|
total_samples = sum(metrics['num_samples'] for metrics in client_metrics) |
|
|
logger.debug(f"Total samples across clients: {total_samples}") |
|
|
for metrics in client_metrics: |
|
|
weight = metrics['num_samples'] / total_samples if self.weighted else 1.0 |
|
|
logger.debug(f"Client metrics: {metrics}, weight: {weight}") |
|
|
for metric_name, value in metrics['metrics'].items(): |
|
|
aggregated_metrics[metric_name] += value * weight |
|
|
logger.info(f"Aggregated metrics: {dict(aggregated_metrics)}") |
|
|
return dict(aggregated_metrics) |
|
|
|
|
|
def check_convergence(self, |
|
|
old_weights: List, |
|
|
new_weights: List, |
|
|
threshold: float = 1e-5) -> bool: |
|
|
logger = logging.getLogger(__name__) |
|
|
logger.debug("Checking convergence...") |
|
|
if old_weights is None or new_weights is None: |
|
|
logger.warning("Old or new weights are None in check_convergence.") |
|
|
return False |
|
|
weight_differences = [ |
|
|
np.mean(np.abs(old - new)) |
|
|
for old, new in zip(old_weights, new_weights) |
|
|
] |
|
|
logger.debug(f"Weight differences: {weight_differences}") |
|
|
converged = all(diff < threshold for diff in weight_differences) |
|
|
logger.info(f"Convergence status: {converged}") |
|
|
return converged |
|
|
|
|
|
|
