How to add new algorithms

Suppose that we are adding the configuration for our new algorithm. New algorithm should be implemented as two classes for server and client. Implementation of the new classes should be derived from the following two base classes:

class appfl.algorithm.BaseServer(weights: OrderedDict, model: torch.nn.Module, loss_fn: torch.nn.Module, num_clients: int, device)

Abstract class of PPFL algorithm for server that aggregates and updates model parameters.

Parameters:

• weight (Dict) – aggregation weight assigned to each client

• model (nn.Module) – torch neural network model to train

• loss_fn (nn.Module) – loss function

• num_clients (int) – the number of clients

• device (str) – device for computation

get_model() → torch.nn.Module

Get the model

Returns:

a deepcopy of self.model

Return type:

nn.Module

class appfl.algorithm.BaseClient(id: int, weight: Dict, model: torch.nn.Module, loss_fn: torch.nn.Module, dataloader: torch.utils.data.DataLoader, cfg, outfile, test_dataloader)

Abstract class of PPFL algorithm for client that trains local model.

Parameters:

• id – unique ID for each client

• weight – aggregation weight assigned to each client

• model – (nn.Module): torch neural network model to train

• loss_fn (nn.Module) – loss function

• dataloader – PyTorch data loader

• device (str) – device for computation

get_model()

Get the model

Returns:

the state_dict of local model

laplace_mechanism_output_perturb(scale_value)

Differential privacy for output perturbation based on Laplacian distribution. This output perturbation adds Laplace noise to primal_state.

Parameters:

scale_value – scaling vector to control the variance of Laplacian distribution

update()

Update local model parameters

Example: NewAlgo

Here we give some simple example.

Core algorithm class

We first create classes for the global and local updates in appfl/algorithm:

• See two classes NewAlgoServer and NewAlgoClient in newalgo.py

• In NewAlgoServer, the update function conducts a global update by averaging the local model parameters sent from multiple clients

• In NewAlgoClient, the update function conducts a local update and send the resulting local model parameters to the server

This is an example code:

Example code for src/appfl/algorithm/newalgo.py

from .algorithm import BaseServer, BaseClient

class NewAlgoServer(BaseServer):
    def __init__(self, weights, model, num_clients, device, **kwargs):
        super(NewAlgoServer, self).__init__(weights, model, num_clients, device)
        self.__dict__.update(kwargs)
        # Any additional initialization

    def update(self, local_states: OrderedDict):
        # Implement new server update function

class NewAlgoClient(BaseClient):
    def __init__(self, id, weight, model, dataloader, device, **kwargs):
        super(NewAlgoClient, self).__init__(id, weight, model, dataloader, device)
        self.__dict__.update(kwargs)
        # Any additional initialization

    def update(self):
        # Implement new client update function

Configuration dataclass

The new algorithm also needs to set up some configurations. This can be done by adding new dataclass under appfl.config.fed. Let’s say we add src/appfl/config/fed/newalgo.py file to implement the dataclass as follows:

Example code for src/appfl/config/fed/newalgo.py

from dataclasses import dataclass
from omegaconf import DictConfig, OmegaConf

@dataclass
class NewAlgo:
    type: str = "newalgo"
    servername: str = "NewAlgoServer"
    clientname: str = "NewAlgoClient"
    args: DictConfig = OmegaConf.create(
        {
            # add new arguments
        }
    )

Then, we need to add the following line to the main configuration file config.py.

from .fed.new_algorithm import *

This is the main configuration class in src/appfl/config/config.py. Each algorithm, specified in Config.fed, can be configured in the dataclasses at appfl.config.fed.*.

The main configuration class

from dataclasses import dataclass, field
from typing import Any
from omegaconf import DictConfig, OmegaConf


from .fed.federated import *
from .fed.iceadmm import *  ## TODO: combine iceadmm and iiadmm under the name of ADMM.
from .fed.iiadmm import *


@dataclass
class Config:
    fed: Any = Federated()

    # Compute device
    device: str = "cpu"

    # Number of training epochs
    num_clients: int = 1

    # Number of training epochs
    num_epochs: int = 2

    # Number of workers in DataLoader
    num_workers: int = 0

    # Train data batch info
    batch_training: bool = True  ## TODO: revisit
    train_data_batch_size: int = 64
    train_data_shuffle: bool = False

    # Indication of whether to validate or not using testing data
    validation: bool = True
    test_data_batch_size: int = 64
    test_data_shuffle: bool = False

    # Checking data sanity
    data_sanity: bool = False

    # Reproducibility
    reproduce: bool = True

    # PCA on Trajectory
    pca_dir: str = ""
    params_start: int=0
    params_end: int=49
    ncomponents: int=40
    
    # Tensorboard
    use_tensorboard: bool = False

    # Loading models
    load_model: bool = False
    load_model_dirname: str = ""
    load_model_filename: str = ""

    # Saving models (server)
    save_model: bool = False
    save_model_dirname: str = ""
    save_model_filename: str = ""
    checkpoints_interval: int = 2

    # Saving state_dict (clients)
    save_model_state_dict: bool = False

    # Logging and recording outputs
    output_dirname: str = "output"
    output_filename: str = "result"
    
    logginginfo: DictConfig = OmegaConf.create({})
    summary_file: str = ""


    #
    # gRPC configutations
    #

    # 100 MB for gRPC maximum message size
    max_message_size: int = 104857600

    operator: DictConfig = OmegaConf.create({"id": 1})
    server: DictConfig = OmegaConf.create(
        {"id": 1, "host": "localhost", "port": 50051, "use_tls": False, "api_key": None}
    )
    client: DictConfig = OmegaConf.create({"id": 1})