Federated Learning으로 다중 공장 설비 고장 예측하기: 데이터 공유 없이 AI 모델 협업 학습 구현

import torch
import torch.nn as nn

class BearingCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv1d(in_channels=3, out_channels=32, kernel_size=5)  # 3축 진동
        self.conv2 = nn.Conv1d(32, 64, kernel_size=5)
        self.pool = nn.MaxPool1d(2)
        self.fc1 = nn.Linear(64 * 248, 128)  # 입력 길이에 따라 조정
        self.fc2 = nn.Linear(128, 1)  # RUL 예측

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = self.pool(x)
        x = torch.relu(self.conv2(x))
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

def local_train(model, train_loader, epochs=5, lr=0.001):
    """
    로컬 데이터로 모델 학습
    """
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()

    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for vibration_data, rul_target in train_loader:
            optimizer.zero_grad()
            output = model(vibration_data)
            loss = criterion(output, rul_target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_loader):.4f}")

    return model.state_dict()  # 학습된 가중치 반환

import copy

def federated_averaging(global_model, client_weights, client_data_sizes):
    """
    FedAvg: 클라이언트 모델을 데이터 크기 가중 평균

    Args:
        global_model: 현재 글로벌 모델
        client_weights: 각 클라이언트의 state_dict 리스트
        client_data_sizes: 각 클라이언트의 데이터 샘플 수 리스트
    """
    total_size = sum(client_data_sizes)
    global_dict = global_model.state_dict()

    # 가중 평균 계산
    for key in global_dict.keys():
        global_dict[key] = torch.zeros_like(global_dict[key])
        for client_weight, size in zip(client_weights, client_data_sizes):
            global_dict[key] += (size / total_size) * client_weight[key]

    global_model.load_state_dict(global_dict)
    return global_model

# 초기화
global_model = BearingCNN()
num_rounds = 20

# 각 공장의 데이터 로더 (실제로는 각 공장 서버에서 실행)
factory_loaders = {
    'A': train_loader_A,  # 10,000 샘플
    'B': train_loader_B,  # 8,000 샘플
    'C': train_loader_C   # 3,000 샘플
}

for round_num in range(num_rounds):
    print(f"\n=== Round {round_num + 1}/{num_rounds} ===")

    client_weights = []
    client_sizes = []

    # 각 클라이언트에서 로컬 학습
    for factory_name, loader in factory_loaders.items():
        print(f"\nFactory {factory_name} training...")

        # 글로벌 모델 복사
        local_model = copy.deepcopy(global_model)

        # 로컬 학습
        weights = local_train(local_model, loader, epochs=5)

        client_weights.append(weights)
        client_sizes.append(len(loader.dataset))

    # 서버에서 집계
    global_model = federated_averaging(global_model, client_weights, client_sizes)
    print(f"\nGlobal model updated (Round {round_num + 1})")

# 최종 모델 저장
torch.save(global_model.state_dict(), 'federated_rul_model.pth')

# Differential Privacy 적용 예시
from opacus import PrivacyEngine

privacy_engine = PrivacyEngine()
model, optimizer, loader = privacy_engine.make_private(
    module=local_model,
    optimizer=optimizer,
    data_loader=train_loader,
    noise_multiplier=1.0,  # 노이즈 강도
    max_grad_norm=1.0      # 그래디언트 클리핑
)

import flwr as fl

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

    def fit(self, parameters, config):
        set_parameters(model, parameters)
        local_train(model, train_loader)
        return self.get_parameters(config), len(train_loader.dataset), {}

    def evaluate(self, parameters, config):
        set_parameters(model, parameters)
        loss, mae = test(model, test_loader)
        return loss, len(test_loader.dataset), {"mae": mae}

fl.client.start_numpy_client(server_address="10.0.0.1:8080", client=BearingClient())