开发: C++知识库 Java知识库 JavaScript Python PHP知识库人工智能区块链大数据移动开发嵌入式开发工具数据结构与算法开发测试游戏开发网络协议系统运维
教程: HTML教程 CSS教程 JavaScript教程 Go语言教程 JQuery教程 VUE教程 VUE3教程 Bootstrap教程 SQL数据库教程 C语言教程 C++教程 Java教程 Python教程 Python3教程 C#教程
数码: 电脑笔记本显卡显示器固态硬盘硬盘耳机手机 iphone vivo oppo 小米华为单反装机图拉丁

-> 人工智能 -> 基于PyTorch实现联邦学习的基本算法FedAvg -> 正文阅读

[人工智能]基于PyTorch实现联邦学习的基本算法FedAvg

I. 前言

在之前的一篇博客联邦学习基本算法FedAvg的代码实现中利用numpy手搭神经网络实现了FedAvg，手搭的神经网络效果已经很好了，不过这还是属于自己造轮子，建议优先使用PyTorch来实现。

II. 数据介绍

联邦学习中存在多个客户端，每个客户端都有自己的数据集，这个数据集他们是不愿意共享的。

本文选用的数据集为中国北方某城市十个区/县从2016年到2019年三年的真实用电负荷数据，采集时间间隔为1小时，即每一天都有24个负荷值。

我们假设这10个地区的电力部门不愿意共享自己的数据，但是他们又想得到一个由所有数据统一训练得到的全局模型。

除了电力负荷数据以外，还有一个备选数据集：风功率数据集。两个数据集通过参数type指定：type == 'load’表示负荷数据，'wind’表示风功率数据。

特征构造

用某一时刻前24个时刻的负荷值以及该时刻的相关气象数据（如温度、湿度、压强等）来预测该时刻的负荷值。

对于风功率数据，同样使用某一时刻前24个时刻的风功率值以及该时刻的相关气象数据来预测该时刻的风功率值。

各个地区应该就如何制定特征集达成一致意见，本文使用的各个地区上的数据的特征是一致的，可以直接使用。

III. 联邦学习

1. 整体框架

原始论文中提出的FedAvg的框架为：
在这里插入图片描述
客户端模型采用PyTorch搭建：

class ANN(nn.Module):
    def __init__(self, input_dim, name, B, E, type, lr):
        super(ANN, self).__init__()
        self.name = name
        self.B = B
        self.E = E
        self.len = 0
        self.type = type
        self.lr = lr
        self.loss = 0
        self.fc1 = nn.Linear(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

2. 服务器端

服务器端执行以下步骤：

初始化参数
对第t轮训练来说：首先计算出 $\cdot K, 1)$ ，然后随机选择m个客户端，对这m个客户端做如下操作（所有客户端并行执行）：更新本地的 $w_t^{k}$ 得到 $w_{t+1}^{k}$ 。所有客户端更新结束后，将 $w_{t+1}^{k}$ 传到服务器，服务器整合所有 $w_{t+1}^{k}$ 得到最新的全局参数 $w_{t+1}$ 。
服务器将最新的 $w_{t+1}$ 分发给所有客户端，然后进行下一轮的更新。

简单来说，每一轮通信时都只是选择部分客户端，这些客户端利用本地的数据进行参数更新，然后将更新后的参数传给服务器，服务器汇总所有客户端的参数形成最新的参数，然后将最新的参数再次分发给所有客户端，进行下一轮更新。

3. 客户端

客户端没什么可说的，就是利用本地数据对神经网络模型的参数进行更新。

IV. 代码实现

1. 初始化

class FedAvg:
    def __init__(self, options):
        self.C = options['C']
        self.E = options['E']
        self.B = options['B']
        self.K = options['K']
        self.r = options['r']
        self.input_dim = options['input_dim']
        self.type = options['type']
        self.lr = options['lr']
        self.clients = options['clients']
        self.nn = ANN(input_dim=self.input_dim, name='server', B=B, E=E, type=self.type, lr=self.lr).to(device)
        self.nns = []
        for i in range(K):
            temp = copy.deepcopy(self.nn)
            temp.name = self.clients[i]
            self.nns.append(temp)

参数：

K，客户端数量，本文为10个，也就是10个地区。
C：选择率，每一轮通信时都只是选择C * K个客户端。
E：客户端更新本地模型的参数时，在本地数据集上训练E轮。
B：客户端更新本地模型的参数时，本地数据集batch大小为B
r：服务器端和客户端一共进行r轮通信。
clients：客户端集合。
type：指定数据类型，负荷预测or风功率预测。
lr：学习率。
input_dim：数据输入维度。
nn：全局模型。
nns：客户端模型集合。

2. 服务器端

服务器端代码如下：

def server(self):
     for t in range(self.r):
          print('第', t + 1, '轮通信:')
          m = np.max([int(self.C * self.K), 1])
          # sampling
          index = random.sample(range(0, self.K), m)
          # local updating
          self.client_update(index)
          # aggregation
          self.aggregation()
          # dispatch
          self.dispatch()

     # return global model
     return self.nn

其中client_update(index)：

def client_update(self, index):  # update nn
     for k in index:
          self.nns[k] = train(self.nns[k])

aggregation():

def aggregation(self):
     s = 0
     for j in range(self.K):
          # normal
          s += self.nns[j].len
          
     params = {}
     with torch.no_grad():
          for k, v in self.nns[0].named_parameters():
               params[k] = copy.deepcopy(v)
               params[k].zero_()
     for j in range(self.K):
          with torch.no_grad():
               for k, v in self.nns[j].named_parameters():
                    params[k] += v * (self.nns[j].len / s)
     with torch.no_grad():
          for k, v in self.nn.named_parameters():
               v.copy_(params[k])

dispatch()：

def dispatch(self):
     params = {}
     with torch.no_grad():
          for k, v in self.nn.named_parameters():
               params[k] = copy.deepcopy(v)
     for j in range(self.K):
          with torch.no_grad():
               for k, v in self.nns[j].named_parameters():
                    v.copy_(params[k])

下面对重要代码进行分析：

客户端的选择

m = np.max([int(self.C * self.K), 1])
index = random.sample(range(0, self.K), m)

index中存储中m个0~10间的整数，表示被选中客户端的序号。

客户端的更新

for k in index:
    self.client_update(self.nns[k])

服务器端汇总客户端模型的参数
$w_{t+1} \gets \sum_{k=1}^{K}\frac{n_k}{n}w_{t+1}^{k}$
其中 $n_k$ 表示第 $k$ 个客户端的本地数据量。也就是说，一个客户端本地的数据越多，它的模型对全局模型的影响就越大。

当然，这只是一种很简单的汇总方式，还有一些其他类型的汇总方式。论文Electricity Consumer Characteristics Identification: A Federated Learning Approach中总结了三种汇总方式：

normal：原始论文中的方式，即根据样本数量来决定客户端参数在最终组合时所占比例。
LA：根据客户端模型的损失占所有客户端损失和的比重来决定最终组合时参数所占比例。
LS：根据损失与样本数量的乘积所占的比重来决定。

值得注意的是，虽然服务器端每次只是选择K个客户端中的m个来进行更新，但在最终汇总的却是所有客户端模型参数。

将更新后的参数分发给客户端

def dispatch(self):
     params = {}
     with torch.no_grad():
          for k, v in self.nn.named_parameters():
               params[k] = copy.deepcopy(v)
     for j in range(self.K):
          with torch.no_grad():
               for k, v in self.nns[j].named_parameters():
                    v.copy_(params[k])

3. 客户端

客户端只需要利用本地数据来进行更新就行了：

def client_update(self, index):  # update nn
     for k in index:
          self.nns[k] = train(self.nns[k])

其中train():

def train(ann):
    ann.train()
    # print(p)
    if ann.type == 'load':
        Dtr, Dte = nn_seq(ann.name, ann.B, ann.type)
    else:
        Dtr, Dte = nn_seq_wind(ann.named, ann.B, ann.type)
    ann.len = len(Dtr)
    # print(len(Dtr))
    loss_function = nn.MSELoss().to(device)
    loss = 0
    optimizer = torch.optim.Adam(ann.parameters(), lr=ann.lr)
    for epoch in range(ann.E):
        cnt = 0
        for (seq, label) in Dtr:
            cnt += 1
            seq = seq.to(device)
            label = label.to(device)
            y_pred = ann(seq)
            loss = loss_function(y_pred, label)
            # add mu/2*|w-wt|**2
            # temp = list(model.parameters())
            # loss += (mu / 2 * () ** 2)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

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

    return ann

4. 测试

def global_test(self):
     model = self.nn
     model.eval()
     c = clients if self.type == 'load' else clients_wind
     for client in c:
          model.name = client
          test(model)

V. 实验及结果

本次实验的参数选择为：

K	C	E	B	r
10	0.5	50	50	5

if __name__ == '__main__':
    K, C, E, B, r = 10, 0.5, 50, 50, 5
    type = 'load'
    input_dim = 30 if type == 'load' else 28
    _client = clients if type == 'load' else clients_wind
    lr = 0.08
    options = {'K': K, 'C': C, 'E': E, 'B': B, 'r': r, 'type': type, 'clients': _client,
               'input_dim': input_dim, 'lr': lr}
    fedavg = FedAvg(options)
    fedavg.server()
    fedavg.global_test()