How to Save Client Models and results After Each Round in Flower about flower

[adam-narozniak]

Hi, You can apply any custom logic in the e.g. fit method of the client. e.g. you can save the results and model state once you're done with fitting.

[kalkite]

Hi @adam-narozniak, thank you for response. I tried to save the client models, but I am only able to save the most last round of communication. Each recent fit file replaces the previous one. However, my goal is to save the client model for each round.

    def fit(self, parameters, config):
        print(f"[Client {int(self.client_id) + 1}] fit, config: {config}")
        self.set_parameters(parameters)
        lr = config['learning_rate']
        optimizer_conf = config['optimizer']
        optimizer = getattr(optim, optimizer_conf)(self.model.parameters(), lr=lr)
        epochs = config['epochs']
        train_loss, train_accuracy = fed_train(model=self.model,
                                               epochs=epochs,
                                               optimizer=optimizer,
                                               train_loader=self.train_loader)
        print(f"Client {int(self.client_id) + 1} train_loss: {train_loss}, train_accuracy: {train_accuracy}")
        # save results of each client to the Results dictionary
        torch.save(self.model.state_dict(), f"client_{int(self.client_id) + 1}_model.pth")
        time.sleep(5)
        return self.get_parameters({}), len(self.train_loader), {}

[adam-narozniak]

Hi @kalkite,
You're close. Now, you can also store the round_id in the config sent. In a simple FedAvg I'd do it as following:

FedAvg(other_params,
              on_fit_config_fn=lambda x: {"round_id": x})

but since you're sending already the e.g. lr it's just one new thing to add in the function to the config.

    def fit(self, parameters, config):
        print(f"[Client {int(self.client_id) + 1}] fit, config: {config}")
        self.set_parameters(parameters)
        lr = config['learning_rate']
        # LINE BELOW IS NEW
        round_id = config['round_id']
        optimizer_conf = config['optimizer']
        optimizer = getattr(optim, optimizer_conf)(self.model.parameters(), lr=lr)
        epochs = config['epochs']
        train_loss, train_accuracy = fed_train(model=self.model,
                                               epochs=epochs,
                                               optimizer=optimizer,
                                               train_loader=self.train_loader)
        print(f"Client {int(self.client_id) + 1} train_loss: {train_loss}, train_accuracy: {train_accuracy}")
        # save results of each client to the Results dictionary
        # MODIFY THE PATH TO INCLUDE THE ROUND ID
        torch.save(self.model.state_dict(), f"client_{int(self.client_id) + 1}_round_{round_id}_model.pth")
        time.sleep(5)
        return self.get_parameters({}), len(self.train_loader), {}

[kalkite]

Is this round_id of the client taken from the server round? Something like this?

def get_on_fit_config(client_configs):
    def fit_config_fn(server_round: int):
        return {
            "learning_rate": client_configs['learning_rate'],
            "optimizer": client_configs['optimizer'],
            "epochs": client_configs['epochs'],
            "round_id": server_round
        }
    return fit_config_fn

I keep on_fit_config_fn = get_on_fit_config(cfg.config_fit)