Mentions légales du service

Skip to content

Draft: declearn-based API revision POC

ANDREY Paul requested to merge poc/declearn into develop

This MR covers a POC to integrate code from declearn as part of the Fed-BioMed backend.

TL;DR:

  • The MAGNET team from Inria has been developing declearn, a python package for SGD-based federated learning, with a focus on modularity and extensibility, notably as to the choice of model framework (scikit-learn, tensorflow, pytorch...) and as to the setup of optimizers on both the client and server sides (using a system of framework-agnostic plug-ins to implement adaptive algorithms, SCAFFOLD, etc.).
  • We (Magnet) believe that some of the design choices and implementation results from declearn could benefit to Fed-BioMed, and would like to merge some of our code and ideas into Fed-BioMed, so that our future efforts may be focused on the latter and/or be easily portable to it.
  • This (Draft) MR covers an initial implementation effort in that direction, and documents additional tasks, ideas and discussions that should either be implemented prior to merging or discussed as part of the planification of the frameworks' convergence.

Important notes:

  • The current effort imports declearn as a third-party package used by the Fed-BioMed code. If the POC is accepted, it will be possible to go further into merging the two frameworks, removing the need for an inter-dependence that may otherwise weaken maintainability effort.
  • The current effort not only affects the API of some backend classes: it also slightly alters the workflow from the researcher end-user's perspective. This POC's review should comprise an in-depth discussion about these alterations, and about more drastic ones that are currently left out of the implementation effort but that I think might be desirable, and easily-achievable thanks to the use of the declearn.
  • The current effort does not cover some functionalities and APIs embarked in declearn, the use and/or inspiration of which could be discussed when planning the future of Fed-BioMed when (if) this POC is accepted, to go further into revising the API and workflow.

Features Tasklist:

  • Revise the TrainingPlan code (and that of its subclasses) to use declearn as backend.
  • Revise the TrainingPlan API sparingly, in the spirit of declearn.
  • Propagate TrainingPlan API changes to other classes (Experiment, Job, Node, Round).
  • Revise the Aggregators API sparingly to use declearn.model.api.Vector, making computations framework-agnostic.
  • Expose the declearn.optimizers.Optimizer API on the researcher side, enabling the use of "FedOpt".
  • Make use of the declearn.optimizers.Optimizer auxiliary-variables API to enable using (declearn-implemented) Scaffold.
  • Restore FedProx, preferably for all frameworks rather than for the sole Torch one.
  • Restore Local and Central Differential Privacy, preferably for all frameworks rather than for the sole Torch one.
  • Add some tensorflow-oriented data-management classes & Test (improve) the already-implemented TensorflowTrainingPlan.

Docs & Tests Tasklist:

  • Revise the 101 example notebook
  • Write a test suite for the revised training plans (possibly use declearn tests as a resource)
  • Revise the tests for altered classes (partially-fixed as of now)
  • Revise non-101 examples and the overall documentation

Bonus tasks (unrelated to the use of declearn):

  • Do some on-the-side cleaning of the revised classes.
  • Clean up import statements and __init__.py files.

Discussion Points:

  • TrainingPlan Workflow
    • keep or roll back currently-implement changes (TrainingPlan being instantiated and wrapping a model)
    • keep or roll back the declearn-backed serialization (that diminishes plan-verification features, but shrinks some attack surface)
  • Data Workflow
    • split up training_data from TrainingPlan (remove unsafe code and burdening import statement strings)
    • either write a data-management class or leave it up to clients
    • enable (de)serializing pre-processing operations (unusable as of now)
  • TrainingArgs and model_args workflow
    • pass everything at plan instantiation time?
    • gather some information from clients, then post_init?
    • avoid passing along information through layers of classes (Experiment->Job->Node->Round)
  • declearn design extension
    • expose tools to extend support to new model (frameworks), optimizer algorithms, etc?
    • revise the integration of declearn.model.api.Model and fedbiomed.common.training_plans.TrainingPlan?
Edited by ANDREY Paul

Merge request reports