연합학습

# Flower 프레임워크를 이용한 연합학습 구현
import flwr as fl
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from collections import OrderedDict

# 간단한 모델 정의
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

# Flower 클라이언트 정의
class FlowerClient(fl.client.NumPyClient):
    def __init__(self, model, train_loader, test_loader):
        self.model = model
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    def get_parameters(self, config):
        """현재 모델 파라미터를 서버로 전송"""
        return [val.cpu().numpy() for val in self.model.state_dict().values()]

    def set_parameters(self, parameters):
        """서버로부터 받은 글로벌 모델로 업데이트"""
        params_dict = zip(self.model.state_dict().keys(), parameters)
        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
        self.model.load_state_dict(state_dict, strict=True)

    def fit(self, parameters, config):
        """로컬 데이터로 학습 (서버 라운드마다 호출)"""
        self.set_parameters(parameters)

        # 로컬 에포크 수 (보통 1-5)
        local_epochs = config.get("local_epochs", 1)

        self.model.train()
        for _ in range(local_epochs):
            for data, target in self.train_loader:
                self.optimizer.zero_grad()
                output = self.model(data)
                loss = self.criterion(output, target)
                loss.backward()
                self.optimizer.step()

        # 업데이트된 파라미터와 데이터 개수 반환
        return self.get_parameters(config), len(self.train_loader.dataset), {}

    def evaluate(self, parameters, config):
        """글로벌 모델 평가"""
        self.set_parameters(parameters)
        self.model.eval()

        loss, correct = 0.0, 0
        with torch.no_grad():
            for data, target in self.test_loader:
                output = self.model(data)
                loss += self.criterion(output, target).item()
                correct += (output.argmax(1) == target).sum().item()

        accuracy = correct / len(self.test_loader.dataset)
        return loss, len(self.test_loader.dataset), {"accuracy": accuracy}

# 클라이언트 실행 (각 디바이스에서)
def start_client(client_id):
    # 각 클라이언트별 로컬 데이터 로드
    train_loader = load_partition(client_id, "train")
    test_loader = load_partition(client_id, "test")

    model = SimpleNet()
    client = FlowerClient(model, train_loader, test_loader)

    # 서버에 연결
    fl.client.start_numpy_client(
        server_address="localhost:8080",
        client=client
    )

# 서버 실행 (중앙 서버에서)
def start_server():
    strategy = fl.server.strategy.FedAvg(
        fraction_fit=0.5,      # 각 라운드에 참여할 클라이언트 비율
        min_fit_clients=2,     # 최소 참여 클라이언트 수
        min_available_clients=3,
    )

    fl.server.start_server(
        server_address="0.0.0.0:8080",
        config=fl.server.ServerConfig(num_rounds=10),  # 10라운드 학습
        strategy=strategy
    )

# 실행: python server.py (서버) / python client.py (각 클라이언트)