Experimenting with Federated learning using PyTorch
Federated learning (FL) is an approach that downloads the current model and computes an updated model at the device itself (ala edge computing) using local data. These locally trained models are then sent from the devices back to the central server where they are aggregated, i.e. averaging weights, and then a single consolidated and improved global model is sent back to the devices.
This project allows you to estimate federated model performance compared to centralize trained on a single device. Emulating clients, federated model being computed by splitting given dataset on equal parts, implementing simple federated averaging weights algorythm. Then single model with same architecture and initial weights learns on whole given dataset. Performance graph of both models on validation dataset is plotted after end of experiment along with saving training logs.
from federated import experiment
experiment(arch, train_ds, val_ds, pref, n_ep, n_clients, n_exp, lr, epoch_step, bs)
arch: clients and main model architecture train_ds: nn.Dataset, training data val_ds: nn.Dataset, validation data pref: prefix to add in experiment result folder name n_ep: total number of training epochs n_clients: number of client models n_exp: number of experiments with different initial weights lr: learning rate epoch_step: number of local training epochs bs: batch size