Simple aggregation schemes on aggregator: geometric mean + median about openfl HOT 3 CLOSED

@alexey-gruzdev, we do have some of this in place: I believe we interpret one of the flplan args as the aggregation function. If the name doesn't match our pre-defined strings, we assume it is a numpy function and invoke it.

I'd like the end-result to be the same as others: you can declare arbitrary python classes in the flplan that provide the functionality, along with init kwargs, the function to invoke, and kwargs for that function. Essentially, our same generic task declaration method we use for the collaborators. This is one of the reasons Patrick and I have been talking about making the aggregator essentially "just another task runner".

This would also have a nice advantage in that we wouldn't have to hard-code the logic that chooses which tensors are inputs for the aggregation function. Instead, it would use the same generic logic that determines the tensor input requirements for any task.

from openfl.

@msheller that's great, if we have it in place - could you provide a simple example for the reference?

Ideally, for better extendibility it would be great to have a python function to be passed to Python API entities & being used by OpenFL further.

def custom_aggregation(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return (a + b)/2

fl_model = fx.run_experiment(collaborators, agg_function=custom_aggregation, 
                                                override_config={'aggregator.settings.rounds_to_train':5})

you can declare arbitrary python classes in the flplan that provide the functionality, along with init kwargs, the function to invoke, and kwargs for that function

Maybe I am wrong, however from Python user's perspective - is it a right place to define such kind of functionalities in FLPlan (.yaml)?

from openfl.

@alexey-gruzdev This is a feature I added a few months ago. A limited set of custom aggregation types are supported at the task level. They can be called in the following way:

plan.yaml
...
tasks:
  aggregated_model_validation:
    aggregation_type: median
    function: validate
    kwargs:
      apply: global
      batch_size: 32
      metrics:
      - accuracy
...

and you can call it through the current Python API like this:

fl_model = fx.run_experiment(collaborators, agg_function=custom_aggregation, 
                                                override_config={'aggregator.settings.rounds_to_train':5,
                                                                 'tasks.aggregated_model_validation.aggregation_type': 'median'})

These aggregation types are currently limited to simple numpy functions (see here). As both of you suggest, we should provide a way to define custom functions that can be applied to model weights and task metrics. There is likely a need to apply a specific aggregation function to certain task output, and possibly another for how to combine model weights. Because of this level of specificity my view is that the plan is the appropriate place to embed this information, but we could extend our existing python APIs so that the user has a cleaner way for expressing the function they want to use without modifying the plan directly.

from openfl.