联邦学习算法FedPer的PyTorch实现

I. 前言

FedPer的原理请见：arXiv | FedPer：带个性化层的联邦学习。

II. 数据集介绍

联邦学习中存在多个客户端，每个客户端都有自己的数据集，这个数据集他们是不愿意共享的。数据集为中国北方某城市十个区/县从2016年到2019年三年的真实用电负荷数据，采集时间间隔为1小时，即每一天都有24个负荷值。

III. FedPer

算法伪代码：

1. 模型定义

客户端的模型为一个简单的四层神经网络模型：

# -*- coding:utf-8 -*-
"""
@Time: 2022/03/03 12:23
@Author: KI
@File: model.py
@Motto: Hungry And Humble
"""
from torch import nn


class ANN(nn.Module):
    def __init__(self, args, name):
        super(ANN, self).__init__()
        self.name = name
        self.len = 0
        self.loss = 0
        self.fc1 = nn.Linear(args.input_dim, 20)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout()
        self.fc2 = nn.Linear(20, 20)
        self.fc3 = nn.Linear(20, 20)
        self.fc4 = nn.Linear(20, 1)

    def forward(self, data):
        x = self.fc1(data)
        x = self.sigmoid(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        x = self.fc3(x)
        x = self.sigmoid(x)
        x = self.fc4(x)
        x = self.sigmoid(x)

        return x

FedPer涉及到基础层和个性化层，四层网络中前三层为为基础层，最后一层为个性化层。

2. 服务器端

服务器端和FedAvg一致，即重复进行客户端采样、参数分发、参数聚合三个步骤。其中参数聚合和参数分发都只针对基础层。

参数分发：

def dispatch(self):
    for j in range(self.args.K):
        cnt = 0
        for old_params, new_params in zip(self.nns[j].parameters(), self.nn.parameters()):
            old_params.data = new_params.data.clone()
            cnt += 1
            if cnt == 2 * (self.args.total - self.args.Kp):
                break

其中Kp为个性化层数量，total为总层数，由于每一层包括权重矩阵和偏置矩阵，所以我们需要乘以2。

同理，参数聚合：

def aggregation(self):
    s = 0
    for j in range(self.args.K):
        # normal
        s += self.nns[j].len

    # 基础层置零
    for v in self.nn.parameters():
        v.data.zero_()

    for j in range(self.args.K):
        cnt = 0
        for v1, v2 in zip(self.nn.parameters(), self.nns[j].parameters()):
            v1.data += v2.data * (self.nns[j].len / s)
            cnt += 1
            if cnt == 2 * (self.args.total - self.args.Kp):
                break

同样是根据客户端样本数量进行聚合。

3. 客户端

客户端训练和FedAvg一致：

def train(args, model, global_round):
    model.train()
    Dtr, Dte, m, n = nn_seq(model.name, args.B)
    model.len = len(Dtr)
    if args.weight_decay != 0:
        lr = args.lr * pow(args.weight_decay, global_round)
    else:
        lr = args.lr
    if args.optimizer == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=lr,
                                     weight_decay=args.weight_decay)
    else:
        optimizer = torch.optim.SGD(model.parameters(), lr=lr,
                                    momentum=0.9)
    print('training...')
    loss_function = nn.MSELoss().to(args.device)
    loss = 0
    for epoch in range(args.E):
        for (seq, label) in Dtr:
            seq = seq.to(args.device)
            label = label.to(args.device)
            y_pred = model(seq)
            loss = loss_function(y_pred, label)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print('epoch', epoch, ':', loss.item())

    return model

IV. 完整代码

后续考虑公开~