Federated Learning API

This module provides tools for implementing federated learning, allowing machine learning models to be trained across multiple decentralized clients without sharing raw data.

FederatedConfig Class

class secureml.federated.FederatedConfig(num_rounds: int = 3, fraction_fit: float = 1.0, min_fit_clients: int = 2, min_available_clients: int = 2, server_address: str = '0.0.0.0:8080', use_secure_aggregation: bool = False, apply_differential_privacy: bool = False, epsilon: float = 1.0, delta: float = 1e-05, weight_update_strategy: str = 'direct', weight_mixing_rate: float = 1.0, weight_momentum: float = 0.9, warmup_rounds: int = 0, apply_weight_constraints: bool = False, max_weight_change: float = 0.2, **kwargs: Any)

Configuration options for federated learning.

__init__(num_rounds: int = 3, fraction_fit: float = 1.0, min_fit_clients: int = 2, min_available_clients: int = 2, server_address: str = '0.0.0.0:8080', use_secure_aggregation: bool = False, apply_differential_privacy: bool = False, epsilon: float = 1.0, delta: float = 1e-05, weight_update_strategy: str = 'direct', weight_mixing_rate: float = 1.0, weight_momentum: float = 0.9, warmup_rounds: int = 0, apply_weight_constraints: bool = False, max_weight_change: float = 0.2, **kwargs: Any)

Initialize federated learning configuration.

Args:: num_rounds: Number of federated training rounds fraction_fit: Fraction of clients used for training in each round min_fit_clients: Minimum number of clients for training min_available_clients: Minimum number of available clients to start round server_address: Server address in the format ‘host:port’ use_secure_aggregation: Whether to use secure aggregation protocol apply_differential_privacy: Whether to apply differential privacy epsilon: Privacy budget for differential privacy (if enabled) delta: Privacy failure probability for differential privacy (if enabled) weight_update_strategy: Strategy for weight updates (“direct”, “ema”, “momentum”) weight_mixing_rate: Weight for new parameters in mixing strategies weight_momentum: Momentum coefficient for momentum strategy warmup_rounds: Number of warmup rounds with gradual mixing rates apply_weight_constraints: Whether to constrain weight updates max_weight_change: Maximum relative change allowed in weights **kwargs: Additional parameters for specific federated learning setups

The FederatedConfig class provides configuration options for federated learning, including parameters for privacy, client selection, and weight update strategies.

Basic Usage Example:

from secureml.federated import FederatedConfig

# Create a configuration for federated learning
config = FederatedConfig(
    num_rounds=5,
    fraction_fit=0.8,
    min_fit_clients=3,
    use_secure_aggregation=True,
    apply_differential_privacy=True,
    epsilon=2.0,
    delta=1e-5
)

Main Functions

secureml.federated.train_federated(model: Any, client_data_fn: Callable[[], Dict[str, DataFrame | ndarray]], config: FederatedConfig | None = None, framework: str = 'auto', model_save_path: str | None = None, **kwargs: Any) → Any

Train a model using federated learning with Flower.

This function sets up a federated learning simulation where the model is trained across multiple clients without centralizing the data.

Args:

model: The model architecture to train (must be compatible with chosen framework) client_data_fn: A function that returns a dictionary mapping client IDs to

their local datasets

config: Configuration for federated learning framework: ML framework to use (‘pytorch’, ‘tensorflow’, or ‘auto’ to detect) model_save_path: Path to save the final federated model **kwargs: Additional parameters passed to client and server setup functions

Returns:

The trained federated model

Raises:

ImportError: If Flower or required ML framework is not installed ValueError: If the framework is not supported or cannot be detected

This function enables training of machine learning models in a federated setting:

from secureml.federated import train_federated
import torch.nn as nn

# Define a model architecture
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 2)
)

# Function to provide client data
def get_client_data():
    return {
        "client1": client1_data,
        "client2": client2_data,
        "client3": client3_data
    }

# Train the model in a federated way
trained_model = train_federated(
    model=model,
    client_data_fn=get_client_data,
    config=config,
    framework="pytorch",
    model_save_path="federated_model.pt",
    batch_size=32,
    epochs=3
)

Server and Client Functions

secureml.federated.start_federated_server(model: Any, config: FederatedConfig | None = None, framework: str = 'auto', **kwargs: Any) → None

Start a Flower federated learning server.

This function starts a server that coordinates the federated learning process among connected clients.

Args:: model: The initial model architecture to distribute config: Configuration for federated learning framework: ML framework to use (‘pytorch’, ‘tensorflow’, or ‘auto’ to detect) **kwargs: Additional parameters for specific server configurations
Raises:: ImportError: If Flower or required ML framework is not installed ValueError: If the framework is not supported or cannot be detected

Start a federated learning server that coordinates model training:

from secureml.federated import start_federated_server, FederatedConfig

# Initialize model
model = create_model()

# Create configuration
config = FederatedConfig(
    server_address="0.0.0.0:8080",
    num_rounds=10,
    min_available_clients=5
)

# Start the server
start_federated_server(
    model=model,
    config=config,
    framework="pytorch"
)

secureml.federated.start_federated_client(model: Any, data: DataFrame | ndarray, server_address: str, framework: str = 'auto', apply_differential_privacy: bool = False, epsilon: float = 1.0, delta: float = 1e-05, **kwargs: Any) → None

Start a Flower federated learning client.

This function starts a client that participates in the federated learning process by training the model on local data and sending the updates to the server.

Args:: model: The model architecture to train locally data: The local training data server_address: Address of the federated learning server (host:port) framework: ML framework to use (‘pytorch’, ‘tensorflow’, or ‘auto’ to detect) apply_differential_privacy: Whether to apply differential privacy to local updates epsilon: Privacy budget for differential privacy (if enabled) delta: Privacy failure probability for differential privacy (if enabled) **kwargs: Additional parameters for specific client configurations
Raises:: ImportError: If Flower or required ML framework is not installed ValueError: If the framework is not supported or cannot be detected

Start a federated learning client that trains the model locally:

from secureml.federated import start_federated_client

# Initialize model with same architecture as server
model = create_model()

# Load local data
local_data = load_local_data()

# Start the client
start_federated_client(
    model=model,
    data=local_data,
    server_address="192.168.1.100:8080",
    framework="pytorch",
    apply_differential_privacy=True,
    epsilon=1.0,
    batch_size=64,
    epochs=2
)

Framework Support

The federated learning module supports both PyTorch and TensorFlow:

PyTorch: For models inheriting from torch.nn.Module
TensorFlow: For models inheriting from tf.keras.Model or tf.Module

When framework="auto" is specified, the framework is detected automatically based on the model type.

Privacy Features

The module supports privacy-preserving techniques:

Secure Aggregation: Protects client model updates using cryptographic techniques
Differential Privacy: Adds calibrated noise to model updates to provide privacy guarantees

Weight Update Strategies

Several weight update strategies are available:

Direct: Standard federated averaging with direct parameter updates
EMA (Exponential Moving Average): Smooth parameter updates using exponential averaging
Momentum: Apply momentum to parameter updates for better convergence

These strategies can be configured using the weight_update_strategy parameter in FederatedConfig.

Best Practices

Test locally first: Use the simulation functionality before deploying to real clients
Start with simpler models: Begin with smaller models before scaling to complex architectures
Monitor privacy budgets: Track epsilon values when using differential privacy
Adjust client parameters: Tune min_fit_clients and fraction_fit based on your client population
Use secure aggregation: Enable use_secure_aggregation in production settings to protect client updates