使用pytorch编写一个使用NSL-KDD数据集训练的联邦学习模型

注：本文为AI（人工智能）语言模型输出的结果，可能需要人工修订。

暂无NSL-KDD数据集，下面给出的是在CIFAR-10数据集上训练的联邦学习模型，可以参考实现方法。

首先，我们需要准备数据集，并按照设定的参与者（client）数量划分数据集。

import copy
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset

# 数据增强以及归一化处理
train_transform = transforms.Compose(
    [
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]
)

test_transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]
)

# 加载CIFAR-10数据集
train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=train_transform)
test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=test_transform)

# 划分成 5 个客户端，并创建对应的 DataLoader
clients_num = 5
clients_data_size = len(train_dataset) // clients_num
clients_datasets = []
for i in range(clients_num):
    start_idx = i * clients_data_size
    end_idx = (i + 1) * clients_data_size if i != clients_num - 1 else len(train_dataset)
    data_subset = Subset(train_dataset, range(start_idx, end_idx))
    clients_datasets.append(data_subset)

clients_data_loaders = [
    DataLoader(dataset, batch_size=64, shuffle=True, drop_last=True) for dataset in clients_datasets
]

test_data_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)

接着，我们需要定义模型以及训练逻辑。

# 定义模型
class CnnClassifier(nn.Module):
    def __init__(self, num_classes=10):
        super(CnnClassifier, self).__init__()

        self.main = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.AvgPool2d(kernel_size=4),
            nn.Flatten(),
            nn.Linear(256, num_classes),
        )

    def forward(self, x):
        return self.main(x)

# 训练单个客户端模型
def train(model, train_loader, criterion, optimizer, device):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        output = model(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * images.shape[0]
        _, predicted = output.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    avg_loss = total_loss / total
    accuracy = correct / total
    return avg_loss, accuracy

# 在测试集上验证多个客户端模型
def validate(model, data_loader, criterion, device):
    model.eval()
    total_loss, correct, total = 0, 0, 0
    with torch.no_grad():
        for images, labels in data_loader:
            images, labels = images.to(device), labels.to(device)
            output = model(images)
            loss = criterion(output, labels)
            total_loss += loss.item() * images.shape[0]
            _, predicted = output.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    avg_loss = total_loss / total
    accuracy = correct / total
    return avg_loss, accuracy

# 定义联邦学习主流程
def federated_train(clients_models, clients_loaders, test_loader, num_epochs, criterion, optimizer, device, model_type="average"):
    """
    clients_models: 初始参与者模型列表
    clients_loaders: 初始参与者训练集 DataLoader 列表
    test_loader: 验证集 DataLoader
    num_epochs: 训练轮数
    criterion: 损失函数
    optimizer: 优化器
    device: 设备
    model_type: 聚合模型类型
        - "last": 取最后一次的模型参数
        - "average": 取平均模型参数
    """
    best_avg_loss = float("inf")
    best_accuracy = 0

    server_model = copy.deepcopy(clients_models[0])
    server_model.to(device)

    for epoch in range(num_epochs):
        print(f"Epoch {epoch + 1}\n{'-' * 10}")

        # 分别在每个客户端上进行训练
        clients_models_copy = copy.deepcopy(clients_models)
        for i, (client_model, train_loader) in enumerate(zip(clients_models_copy, clients_loaders)):
            print(f"Client {i}")
            client_loss, client_acc = train(client_model, train_loader, criterion, optimizer, device)
            print(f"Loss: {client_loss:.4f} | Accuracy: {client_acc:.4f}\n")

        # 使用“平均”或“最后”的方式聚合所有参与者的模型参数
        if model_type == "average":
            for param in server_model.parameters():
                param.data *= 0

            for client_model in clients_models_copy:
                for param, client_param in zip(server_model.parameters(), client_model.parameters()):
                    param.data += client_param.data
            for param in server_model.parameters():
                param.data /= len(clients_models)

        elif model_type == "last":
            server_model = copy.deepcopy(clients_models_copy[-1])

        # 在验证集上验证聚合后的模型
        test_loss, test_acc = validate(server_model, test_loader, criterion, device)
        print(f"Test Loss: {test_loss:.4f} | Test Accuracy: {test_acc:.4f}\n")

        # 根据模型准确率进行模型保存
        if test_loss <= best_avg_loss and test_acc >= best_accuracy:
            print(f"Saving model...\n")
            best_avg_loss = test_loss
            best_accuracy = test_acc

    return server_model, best_avg_loss, best_accuracy

接着，我们可以初始化参与者模型并运行联邦学习。

# 初始化参与者模型
clients_models = [CnnClassifier() for _ in range(clients_num)]

# 定义超参数
num_epochs = 10
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(clients_models[0].parameters(), lr=0.1, momentum=0.9)

# 开始联邦学习
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
federated_train(clients_models, clients_data_loaders, test_data_loader, num_epochs, criterion,
                optimizer, device, model_type="average")