PFLlib

Easy to Extend

This library is designed to be easily extendable with new algorithms and datasets. Here’s how you can add them:

• New Dataset: To add a new dataset, simply create a generate_DATA.py file in ./dataset and then write the download code and use the utils as shown in ./dataset/generate_MNIST.py (you can consider it as a template):

  # `generate_DATA.py`
  import necessary pkgs
  from utils import necessary processing funcs
  
  def generate_dataset(...):
      # download dataset as usual
      # pre-process dataset as usual
      X, y, statistic = separate_data((dataset_content, dataset_label), ...)
      train_data, test_data = split_data(X, y)
      save_file(config_path, train_path, test_path, train_data, test_data, statistic, ...)
  
  # call the generate_dataset func

• New Algorithm: To add a new algorithm, extend the base classes Server and Client, which are defined in ./system/flcore/servers/serverbase.py and ./system/flcore/clients/clientbase.py, respectively.
  • Server

    # serverNAME.py
    import necessary pkgs
    from flcore.clients.clientNAME import clientNAME
    from flcore.servers.serverbase import Server
    
    class NAME(Server):
        def __init__(self, args, times):
            super().__init__(args, times)
    
            # select slow clients
            self.set_slow_clients()
            self.set_clients(clientAVG)
        def train(self):
            # server scheduling code of your algorithm

  • Client

    # clientNAME.py
    import necessary pkgs
    from flcore.clients.clientbase import Client
    
    class clientNAME(Client):
        def __init__(self, args, id, train_samples, test_samples, **kwargs):
            super().__init__(args, id, train_samples, test_samples, **kwargs)
            # add specific initialization
        
        def train(self):
            # client training code of your algorithm

• New Model: To add a new model, simply include it in ./system/flcore/trainmodel/models.py.
• New Optimizer: If you need a new optimizer for training, add it to ./system/flcore/optimizers/fedoptimizer.py.
• New Benchmark Platform or Library: Our code framework is flexible, allowing users to build custom platforms or libraries for specific applications, such as FL-IoT and HtFLlib.