osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
537 lines
17 KiB
Rust
537 lines
17 KiB
Rust
//! FederatedMed - Privacy-Preserving Medical AI.
|
|
//!
|
|
//! This demo implements federated learning for training on hospital data
|
|
//! without sharing sensitive patient information.
|
|
|
|
pub mod aggregator;
|
|
pub mod privacy;
|
|
pub mod sample_data;
|
|
|
|
use rand::Rng;
|
|
use thiserror::Error;
|
|
|
|
use fedmed_shared::{
|
|
AggregationStrategy, ClientInfo, FederatedConfig, FederationResult, GlobalModel, LocalUpdate,
|
|
PrivacyConfig, TestMetrics, TrainingPhase, TrainingProgress,
|
|
};
|
|
|
|
use aggregator::Aggregator;
|
|
use privacy::{DifferentialPrivacy, GradientClipper, PrivacyAccountant};
|
|
|
|
/// Errors that can occur in FederatedMed.
|
|
#[derive(Debug, Error)]
|
|
pub enum FederatedError {
|
|
/// Invalid configuration.
|
|
#[error("Invalid configuration: {0}")]
|
|
InvalidConfig(String),
|
|
|
|
/// Training failed.
|
|
#[error("Training failed: {0}")]
|
|
TrainingFailed(String),
|
|
|
|
/// Not enough clients.
|
|
#[error("Not enough clients: need {needed}, have {available}")]
|
|
NotEnoughClients { needed: usize, available: usize },
|
|
|
|
/// Privacy budget exhausted.
|
|
#[error("Privacy budget exhausted: spent {spent}, limit {limit}")]
|
|
PrivacyBudgetExhausted { spent: f64, limit: f64 },
|
|
|
|
/// Aggregation failed.
|
|
#[error("Aggregation failed: {0}")]
|
|
AggregationFailed(String),
|
|
|
|
/// Client error.
|
|
#[error("Client error: {0}")]
|
|
ClientError(String),
|
|
}
|
|
|
|
/// Main Federated Learning Coordinator.
|
|
#[derive(Debug)]
|
|
pub struct FederatedCoordinator {
|
|
/// Configuration for federated learning.
|
|
config: FederatedConfig,
|
|
/// Global model.
|
|
global_model: GlobalModel,
|
|
/// Registered clients.
|
|
clients: Vec<ClientInfo>,
|
|
/// Aggregator for combining updates.
|
|
aggregator: Aggregator,
|
|
/// Privacy mechanism.
|
|
privacy: Option<DifferentialPrivacy>,
|
|
/// Privacy accountant for budget tracking.
|
|
accountant: Option<PrivacyAccountant>,
|
|
/// Gradient clipper.
|
|
clipper: Option<GradientClipper>,
|
|
/// Current round.
|
|
current_round: u64,
|
|
/// Training history (accuracy per round).
|
|
accuracy_history: Vec<f64>,
|
|
/// Training history (loss per round).
|
|
loss_history: Vec<f64>,
|
|
/// Total privacy budget spent.
|
|
privacy_spent: f64,
|
|
}
|
|
|
|
impl Default for FederatedCoordinator {
|
|
fn default() -> Self {
|
|
Self::new(FederatedConfig::default())
|
|
}
|
|
}
|
|
|
|
impl FederatedCoordinator {
|
|
/// Create a new federated coordinator.
|
|
#[must_use]
|
|
pub fn new(config: FederatedConfig) -> Self {
|
|
let aggregator = Aggregator::new(config.aggregation_strategy);
|
|
|
|
let (privacy, accountant, clipper) = if let Some(ref privacy_config) = config.privacy {
|
|
(
|
|
Some(DifferentialPrivacy::new(privacy_config.clone())),
|
|
Some(PrivacyAccountant::new(privacy_config.clone())),
|
|
Some(GradientClipper::new(privacy_config.clip_norm)),
|
|
)
|
|
} else {
|
|
(None, None, None)
|
|
};
|
|
|
|
Self {
|
|
config,
|
|
global_model: GlobalModel::default(),
|
|
clients: Vec::new(),
|
|
aggregator,
|
|
privacy,
|
|
accountant,
|
|
clipper,
|
|
current_round: 0,
|
|
accuracy_history: Vec::new(),
|
|
loss_history: Vec::new(),
|
|
privacy_spent: 0.0,
|
|
}
|
|
}
|
|
|
|
/// Register a client with the federation.
|
|
pub fn register_client(&mut self, client: ClientInfo) {
|
|
self.clients.push(client);
|
|
}
|
|
|
|
/// Register multiple clients.
|
|
pub fn register_clients(&mut self, clients: Vec<ClientInfo>) {
|
|
self.clients.extend(clients);
|
|
}
|
|
|
|
/// Get the number of registered clients.
|
|
#[must_use]
|
|
pub fn num_clients(&self) -> usize {
|
|
self.clients.len()
|
|
}
|
|
|
|
/// Get active clients.
|
|
#[must_use]
|
|
pub fn active_clients(&self) -> Vec<&ClientInfo> {
|
|
self.clients.iter().filter(|c| c.is_active).collect()
|
|
}
|
|
|
|
/// Select clients for a round of training.
|
|
fn select_clients(&self) -> Vec<&ClientInfo> {
|
|
let active: Vec<_> = self.active_clients();
|
|
let num_to_select =
|
|
((active.len() as f32 * self.config.client_fraction).ceil() as usize).max(1);
|
|
|
|
let mut rng = rand::thread_rng();
|
|
let mut indices: Vec<usize> = (0..active.len()).collect();
|
|
|
|
// Fisher-Yates shuffle for random selection
|
|
for i in (1..indices.len()).rev() {
|
|
let j = rng.gen_range(0..=i);
|
|
indices.swap(i, j);
|
|
}
|
|
|
|
indices
|
|
.into_iter()
|
|
.take(num_to_select)
|
|
.map(|i| active[i])
|
|
.collect()
|
|
}
|
|
|
|
/// Simulate local training at a client.
|
|
fn simulate_local_training(&self, client: &ClientInfo) -> LocalUpdate {
|
|
let mut rng = rand::thread_rng();
|
|
|
|
// Simulate weight updates
|
|
let num_params = self.global_model.num_parameters;
|
|
let weight_deltas: Vec<f32> = (0..num_params)
|
|
.map(|_| rng.gen_range(-0.1_f32..0.1_f32))
|
|
.collect();
|
|
|
|
// Simulate local metrics
|
|
let base_accuracy = 0.7 + (client.data_size as f64 / 10000.0).min(0.2);
|
|
let noise: f64 = rng.gen_range(-0.05..0.05);
|
|
|
|
LocalUpdate {
|
|
client_id: client.id.clone(),
|
|
round: self.current_round,
|
|
weight_deltas,
|
|
num_samples: client.data_size,
|
|
local_loss: 0.5 - (self.current_round as f64 * 0.005),
|
|
local_accuracy: base_accuracy + noise + (self.current_round as f64 * 0.002),
|
|
training_time: (client.data_size as f64 / 100.0) + rng.gen_range(0.0..5.0),
|
|
control_variates: None,
|
|
}
|
|
}
|
|
|
|
/// Run one federated round.
|
|
pub fn run_round(&mut self) -> Result<TrainingProgress, FederatedError> {
|
|
// Check if we have enough clients
|
|
let active_count = self.active_clients().len();
|
|
if active_count < self.config.min_clients {
|
|
return Err(FederatedError::NotEnoughClients {
|
|
needed: self.config.min_clients,
|
|
available: active_count,
|
|
});
|
|
}
|
|
|
|
// Check privacy budget
|
|
if let Some(ref privacy_config) = self.config.privacy
|
|
&& let Some(target) = privacy_config.target_epsilon
|
|
&& self.privacy_spent >= target
|
|
{
|
|
return Err(FederatedError::PrivacyBudgetExhausted {
|
|
spent: self.privacy_spent,
|
|
limit: target,
|
|
});
|
|
}
|
|
|
|
self.current_round += 1;
|
|
|
|
// Select clients
|
|
let selected_clients = self.select_clients();
|
|
let num_participating = selected_clients.len();
|
|
|
|
// Collect local updates
|
|
let mut updates: Vec<LocalUpdate> = selected_clients
|
|
.iter()
|
|
.map(|client| {
|
|
let mut update = self.simulate_local_training(client);
|
|
|
|
// Apply gradient clipping if enabled
|
|
if let Some(ref clipper) = self.clipper {
|
|
update.weight_deltas = clipper.clip(&update.weight_deltas);
|
|
}
|
|
|
|
// Add noise if differential privacy is enabled
|
|
if let Some(ref dp) = self.privacy {
|
|
update.weight_deltas = dp.add_noise(&update.weight_deltas);
|
|
}
|
|
|
|
update
|
|
})
|
|
.collect();
|
|
|
|
// Track privacy budget
|
|
if let Some(ref mut accountant) = self.accountant {
|
|
let budget_spent = accountant.step(updates.len());
|
|
self.privacy_spent += budget_spent;
|
|
}
|
|
|
|
// Aggregate updates
|
|
let aggregated = self
|
|
.aggregator
|
|
.aggregate(&mut updates, &self.global_model)?;
|
|
|
|
// Update global model
|
|
self.update_global_model(&aggregated);
|
|
|
|
// Evaluate
|
|
let (accuracy, loss) = self.evaluate();
|
|
self.accuracy_history.push(accuracy);
|
|
self.loss_history.push(loss);
|
|
|
|
// Update global model metrics
|
|
self.global_model.validation_accuracy = Some(accuracy);
|
|
self.global_model.validation_loss = Some(loss);
|
|
|
|
Ok(TrainingProgress {
|
|
round: self.current_round,
|
|
total_rounds: self.config.num_rounds as u64,
|
|
accuracy,
|
|
loss,
|
|
privacy_budget_spent: self.privacy_spent,
|
|
participating_clients: num_participating,
|
|
total_clients: self.clients.len(),
|
|
elapsed_seconds: self.current_round as f64 * 60.0, // Simulated
|
|
eta_seconds: Some((self.config.num_rounds as u64 - self.current_round) as f64 * 60.0),
|
|
phase: TrainingPhase::Aggregation,
|
|
})
|
|
}
|
|
|
|
/// Update global model with aggregated weights.
|
|
fn update_global_model(&mut self, aggregated_deltas: &[f32]) {
|
|
for (weight, delta) in self
|
|
.global_model
|
|
.weights
|
|
.iter_mut()
|
|
.zip(aggregated_deltas.iter())
|
|
{
|
|
*weight += delta * self.config.learning_rate as f32;
|
|
}
|
|
self.global_model.version = self.current_round;
|
|
}
|
|
|
|
/// Evaluate the global model.
|
|
fn evaluate(&self) -> (f64, f64) {
|
|
// Simulated evaluation
|
|
let base_accuracy = 0.5;
|
|
let improvement = (self.current_round as f64 * 0.008).min(0.45);
|
|
let accuracy = base_accuracy + improvement;
|
|
|
|
let base_loss = 2.0;
|
|
let loss_reduction = (self.current_round as f64 * 0.03).min(1.5);
|
|
let loss = base_loss - loss_reduction;
|
|
|
|
(accuracy, loss)
|
|
}
|
|
|
|
/// Run complete federated training.
|
|
pub fn train(
|
|
&mut self,
|
|
progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
|
|
) -> Result<FederationResult, FederatedError> {
|
|
// Validate configuration
|
|
if self.clients.is_empty() {
|
|
return Err(FederatedError::InvalidConfig(
|
|
"No clients registered".to_string(),
|
|
));
|
|
}
|
|
|
|
for round in 0..self.config.num_rounds {
|
|
let progress = self.run_round()?;
|
|
|
|
if let Some(ref callback) = progress_callback {
|
|
callback(progress);
|
|
}
|
|
|
|
// Early stopping check
|
|
if round > 10 && self.accuracy_history.len() >= 5 {
|
|
let recent: Vec<_> = self.accuracy_history.iter().rev().take(5).collect();
|
|
let improvement = recent[0] - recent[4];
|
|
if improvement < 0.001 {
|
|
break; // Convergence
|
|
}
|
|
}
|
|
}
|
|
|
|
// Calculate final test metrics
|
|
let test_metrics = self.calculate_test_metrics();
|
|
|
|
// Calculate client statistics
|
|
let client_stats = self.calculate_client_stats();
|
|
|
|
Ok(FederationResult {
|
|
final_model: self.global_model.clone(),
|
|
accuracy_history: self.accuracy_history.clone(),
|
|
loss_history: self.loss_history.clone(),
|
|
total_time: self.current_round as f64 * 60.0,
|
|
rounds_completed: self.current_round,
|
|
total_privacy_budget: self.privacy_spent,
|
|
client_stats,
|
|
test_metrics,
|
|
})
|
|
}
|
|
|
|
/// Calculate test metrics.
|
|
fn calculate_test_metrics(&self) -> TestMetrics {
|
|
let accuracy = self.accuracy_history.last().copied().unwrap_or(0.0);
|
|
|
|
TestMetrics {
|
|
accuracy,
|
|
precision: accuracy * 0.98,
|
|
recall: accuracy * 0.96,
|
|
f1_score: accuracy * 0.97,
|
|
auc_roc: accuracy * 1.02, // AUC typically slightly higher
|
|
confusion_matrix: vec![85, 5, 3, 7], // Example 2-class confusion matrix
|
|
num_classes: 2,
|
|
}
|
|
}
|
|
|
|
/// Calculate client statistics.
|
|
fn calculate_client_stats(&self) -> Vec<fedmed_shared::ClientStats> {
|
|
self.clients
|
|
.iter()
|
|
.map(|client| {
|
|
let rounds_participated = (self.current_round as f64
|
|
* self.config.client_fraction as f64)
|
|
.ceil() as usize;
|
|
|
|
fedmed_shared::ClientStats {
|
|
client_id: client.id.clone(),
|
|
rounds_participated,
|
|
avg_local_accuracy: 0.85,
|
|
total_samples: client.data_size * rounds_participated,
|
|
avg_training_time: client.data_size as f64 / 100.0,
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Get current global model.
|
|
#[must_use]
|
|
pub fn global_model(&self) -> &GlobalModel {
|
|
&self.global_model
|
|
}
|
|
|
|
/// Get current round.
|
|
#[must_use]
|
|
pub fn current_round(&self) -> u64 {
|
|
self.current_round
|
|
}
|
|
|
|
/// Get privacy budget spent.
|
|
#[must_use]
|
|
pub fn privacy_spent(&self) -> f64 {
|
|
self.privacy_spent
|
|
}
|
|
|
|
/// Get accuracy history.
|
|
#[must_use]
|
|
pub fn accuracy_history(&self) -> &[f64] {
|
|
&self.accuracy_history
|
|
}
|
|
|
|
/// Get loss history.
|
|
#[must_use]
|
|
pub fn loss_history(&self) -> &[f64] {
|
|
&self.loss_history
|
|
}
|
|
|
|
/// Get configuration.
|
|
#[must_use]
|
|
pub fn config(&self) -> &FederatedConfig {
|
|
&self.config
|
|
}
|
|
|
|
/// Update aggregation strategy.
|
|
pub fn set_aggregation_strategy(&mut self, strategy: AggregationStrategy) {
|
|
self.aggregator = Aggregator::new(strategy);
|
|
}
|
|
|
|
/// Update privacy configuration.
|
|
pub fn set_privacy_config(&mut self, config: PrivacyConfig) {
|
|
self.privacy = Some(DifferentialPrivacy::new(config.clone()));
|
|
self.accountant = Some(PrivacyAccountant::new(config.clone()));
|
|
self.clipper = Some(GradientClipper::new(config.clip_norm));
|
|
}
|
|
}
|
|
|
|
/// Run the demo.
|
|
pub fn run_demo() -> Result<FederationResult, FederatedError> {
|
|
// Create coordinator with sample config
|
|
let config = sample_data::hospital_network_config();
|
|
let mut coordinator = FederatedCoordinator::new(config);
|
|
|
|
// Register sample clients
|
|
let clients = sample_data::chest_xray_federation();
|
|
coordinator.register_clients(clients);
|
|
|
|
// Run training
|
|
coordinator.train(None)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_coordinator_creation() {
|
|
let coordinator = FederatedCoordinator::default();
|
|
assert_eq!(coordinator.current_round(), 0);
|
|
assert_eq!(coordinator.num_clients(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_register_clients() {
|
|
let mut coordinator = FederatedCoordinator::default();
|
|
let clients = fedmed_shared::sample_clients();
|
|
coordinator.register_clients(clients);
|
|
assert_eq!(coordinator.num_clients(), 5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_run_round() {
|
|
let mut coordinator = FederatedCoordinator::default();
|
|
coordinator.register_clients(fedmed_shared::sample_clients());
|
|
|
|
let result = coordinator.run_round();
|
|
assert!(result.is_ok());
|
|
|
|
let progress = result.unwrap();
|
|
assert_eq!(progress.round, 1);
|
|
assert!(progress.participating_clients > 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_not_enough_clients() {
|
|
let config = FederatedConfig {
|
|
min_clients: 10,
|
|
..Default::default()
|
|
};
|
|
let mut coordinator = FederatedCoordinator::new(config);
|
|
coordinator.register_clients(vec![ClientInfo::default()]);
|
|
|
|
let result = coordinator.run_round();
|
|
assert!(matches!(
|
|
result,
|
|
Err(FederatedError::NotEnoughClients { .. })
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn test_training() {
|
|
let config = FederatedConfig {
|
|
num_rounds: 5,
|
|
min_clients: 2,
|
|
privacy: None, // Disable privacy for this test to avoid budget exhaustion
|
|
..Default::default()
|
|
};
|
|
let mut coordinator = FederatedCoordinator::new(config);
|
|
coordinator.register_clients(fedmed_shared::sample_clients());
|
|
|
|
let result = coordinator.train(None);
|
|
assert!(result.is_ok());
|
|
|
|
let federation_result = result.unwrap();
|
|
assert!(federation_result.rounds_completed > 0);
|
|
assert!(!federation_result.accuracy_history.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_privacy_tracking() {
|
|
let config = FederatedConfig {
|
|
num_rounds: 3,
|
|
privacy: Some(fedmed_shared::sample_privacy_config()),
|
|
..Default::default()
|
|
};
|
|
let mut coordinator = FederatedCoordinator::new(config);
|
|
coordinator.register_clients(fedmed_shared::sample_clients());
|
|
|
|
let _ = coordinator.run_round();
|
|
assert!(coordinator.privacy_spent() > 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_set_aggregation_strategy() {
|
|
let mut coordinator = FederatedCoordinator::default();
|
|
coordinator.set_aggregation_strategy(AggregationStrategy::FedProx);
|
|
// Strategy is set internally
|
|
}
|
|
|
|
#[test]
|
|
fn test_run_demo() {
|
|
let result = run_demo();
|
|
if let Err(ref e) = result {
|
|
eprintln!("run_demo error: {:?}", e);
|
|
}
|
|
assert!(result.is_ok());
|
|
}
|
|
}
|