Closed jeandut closed 2 years ago
As some parts of the code are common to all strategies (e.g., training/evaluation loops, computing the average of models), may be we should factorize those elements. I suggest that we have a file utils.py
in FLamby/flamby/strategies
with all training utilities. What do you think?
Thanks @omarfoq ! You are not wrong theoretically. In practice as we are working with lots of people asynchronously it might be hard to achieve a perfect synergy. Let's see what people come up with for the first strategies and we'll see how much we can factorize the code. Note that for evaluation there is already a util.
FedAvg is now implemented in the repository. There is a utils.py
file in the strategies folder where we define two abstractions that might or not be useful for your strategies as well as @omarfoq hinted at.
FLamby/flamby/strategies
strategy_name.py
(aka fed_avg.py/fed_yogi.py, scaffold.y). This file should contain a classStrategyName
with the following methods and init:class StrategyName(): def init(self, training_dataloaders: list[torch.dataset], model: torch.nn.Module, loss: torch.nn.modules.loss._Loss, learning_rate: float, nrounds: int, additional_strategy_specific_parameters): self.training_dataloaders = training_dataloaders self.modelslist = [copy.deepcopy(model) for in range(len(training_dataloaders))] self.loss = loss self.lr = learning_rate self.nrounds = nrounds def perform_round(self):
do stuff and update models