client.py

from collections import OrderedDict
from centralized import load_data, load_model, train, test

import flwr as fl
import torch


def set_parameters(model, parameters):
    params_dict = zip(model.state_dict().keys(), parameters)
    state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
    model.load_state_dict(state_dict, strict=True)

net = load_model()
trainloader, testloader = load_data()

class FlowerClient(fl.client.NumPyClient):
    def get_parameters(self, config):
        return [val.cpu().numpy() for _, val in net.state_dict().items()]

    def fit(self, parameters, config):
        set_parameters(net, parameters)
        train(net, trainloader, epoch=1)
        return self.get_parameters({}), len(trainloader.dataset), {}

    def evaluate(self, parameters, config):
        set_parameters(net, parameters)
        loss, accuracy = test(net, testloader)
        print(f"Evaluation: Loss = {loss}, Accuracy = {accuracy}")
        return float(loss), len(testloader.dataset), {"accuracy": float(accuracy)}


fl.client.start_client(
    server_address="127.0.0.1:8080",
    client=FlowerClient().to_client(),
)