"""Federated Learning Client for Myanmar Ghost project. Enables training on distributed data (e.g., hospitals, restaurants) without centralizing sensitive data. """ import flwr as fl import numpy as np import torch import torch.nn as nn import torch.optim as optim from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import yaml logger = __import__("loguru").logger class FederatedClient(fl.client.NumPyClient): """Flower client for federated learning.""" def __init__( self, model: nn.Module, trainloader, valloader, device: str = "cuda" if torch.cuda.is_available() else "cpu", client_id: str = "client_1", output_dir: str = "outputs/federated", ): self.model = model.to(device) self.trainloader = trainloader self.valloader = valloader self.device = device self.client_id = client_id self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.criterion = nn.CrossEntropyLoss() self.optimizer = optim.AdamW(self.model.parameters(), lr=1e-4) def get_parameters(self) -> List[np.ndarray]: """Get model parameters as numpy arrays.""" return [val.cpu().numpy() for _, val in self.model.state_dict().items()] def set_parameters(self, parameters: List[np.ndarray]) -> None: """Set model parameters from numpy arrays.""" state_dict = dict(self.model.state_dict()) for i, (key, _) in enumerate(state_dict.items()): state_dict[key] = torch.from_numpy(parameters[i]) self.model.load_state_dict(state_dict) def fit( self, parameters: List[np.ndarray], config: Dict[str, Any], ) -> Tuple[List[np.ndarray], int, Dict]: """Train model on local data. Args: parameters: Global model parameters config: Training configuration Returns: Updated parameters, number of samples, metrics """ # Set global parameters self.set_parameters(parameters) # Training configuration epochs = config.get("local_epochs", 1) batch_size = config.get("batch_size", 32) learning_rate = config.get("learning_rate", 1e-4) self.optimizer = optim.AdamW( self.model.parameters(), lr=learning_rate, ) # Local training self.model.train() total_loss = 0.0 total_samples = 0 for epoch in range(epochs): epoch_loss = 0.0 epoch_samples = 0 for batch_idx, (inputs, labels) in enumerate(self.trainloader): inputs = inputs.to(self.device) labels = labels.to(self.device) self.optimizer.zero_grad() outputs = self.model(inputs) loss = self.criterion(outputs, labels) loss.backward() self.optimizer.step() epoch_loss += loss.item() * inputs.size(0) epoch_samples += inputs.size(0) total_loss += epoch_loss total_samples += epoch_samples logger.info( f"Client {self.client_id} - Epoch {epoch+1}/{epochs}: " f"Loss={epoch_loss/epoch_samples:.4f}" ) # Save checkpoint self._save_checkpoint(epochs) metrics = { "loss": total_loss / total_samples, "samples": total_samples, "epochs": epochs, } return self.get_parameters(), total_samples, metrics def evaluate( self, parameters: List[np.ndarray], config: Dict[str, Any], ) -> Tuple[float, int, Dict]: """Evaluate model on local validation data. Args: parameters: Model parameters config: Evaluation configuration Returns: Loss, number of samples, metrics """ self.set_parameters(parameters) self.model.eval() total_loss = 0.0 total_correct = 0 total_samples = 0 with torch.no_grad(): for inputs, labels in self.valloader: inputs = inputs.to(self.device) labels = labels.to(self.device) outputs = self.model(inputs) loss = self.criterion(outputs, labels) total_loss += loss.item() * inputs.size(0) _, predicted = outputs.max(1) total_correct += predicted.eq(labels).sum().item() total_samples += inputs.size(0) accuracy = total_correct / total_samples if total_samples > 0 else 0.0 logger.info( f"Client {self.client_id} - Evaluation: " f"Loss={total_loss/total_samples:.4f}, Accuracy={accuracy:.4f}" ) metrics = { "loss": total_loss / total_samples, "accuracy": accuracy, "samples": total_samples, } return total_loss / total_samples, total_samples, metrics def _save_checkpoint(self, epochs: int) -> None: """Save model checkpoint.""" path = self.output_dir / f"{self.client_id}_checkpoint.pt" torch.save({ "model_state_dict": self.model.state_dict(), "optimizer_state_dict": self.optimizer.state_dict(), "epochs": epochs, "client_id": self.client_id, }, path) logger.info(f"Checkpoint saved to {path}") def load_client_config(config_path: str) -> Dict: """Load client configuration from YAML.""" with open(config_path, "r", encoding="utf-8") as f: return yaml.safe_load(f) class ClientFactory: """Factory for creating federated clients.""" def __init__( self, model_fn, data_dir: str, output_dir: str = "outputs/federated", ): self.model_fn = model_fn self.data_dir = Path(data_dir) self.output_dir = Path(output_dir) def create_client( self, client_id: str, config: Dict, ) -> FederatedClient: """Create a client for the given configuration.""" from torch.utils.data import DataLoader # Load data train_data = self._load_partition( client_id, config.get("train_file", f"{client_id}_train.pt"), ) val_data = self._load_partition( client_id, config.get("val_file", f"{client_id}_val.pt"), ) trainloader = DataLoader( train_data, batch_size=config.get("batch_size", 32), shuffle=True, ) valloader = DataLoader( val_data, batch_size=config.get("batch_size", 32), shuffle=False, ) model = self.model_fn() return FederatedClient( model=model, trainloader=trainloader, valloader=valloader, device=config.get("device", "cuda"), client_id=client_id, output_dir=str(self.output_dir), ) def _load_partition(self, client_id: str, filename: str): """Load data partition for a client.""" path = self.data_dir / client_id / filename if path.exists(): return torch.load(path) raise FileNotFoundError(f"Data partition not found: {path}") def start_client( model_fn, data_dir: str, client_id: str, server_address: str = "localhost:8080", config_path: Optional[str] = None, ) -> None: """Start a federated learning client.""" config = {} if config_path: config = load_client_config(config_path) factory = ClientFactory(model_fn, data_dir) client = factory.create_client(client_id, config) app = fl.client.start_numpy_client( server_address=server_address, client=client, ) if __name__ == "__main__": print("FederatedClient module loaded") print("Use start_client() to start a federated learning client")