Pysyft学习笔记三：分布式模型实现（Fed_avg算法整合模型）

Posted 2022-02-08 一只特立独行的猫

导入包

导入包，并初始化Hook

import torch 
import syft as sy

from torch import nn
from torch import optim

hook = sy.TorchHook(torch)

构建客户机

Bob = sy.VirtualWorker(hook,id='Bob')
Alice = sy.VirtualWorker(hook,id='Alice')

构建数据集

自己随便想了一些数据来训练

data = torch.tensor([[0,0],[0,1], [1,0], [1,1.]],requires_grad=True)
targe = torch.tensor([[0],[0],[1], [1.]],requires_grad=True)

构建模型

先建立一层全连接层，然后把数据分发出去，相当于初始化了训练环境。

#全连接模型，输入是二维，输出是一维，可以理解为读两个数，输出一个数
model = nn.Linear(2,1)

data_Bob_ptr = data[:2].send(Bob)
targe_Bob_ptr = targe[:2].send(Bob)
data_Alice_ptr = data[2:].send(Alice)
targe_Alice_ptr = data[2:].send(Alice)

定义分布式训练函数，采用平均算法合并模型

想把初始模型发送到客户端，客户端用自己所拥有的数据进行训练后，再将模型返回给服务器。服务器将两个模型进行整合，得到一个对双方数据都有较好预测效果的模型。

def train():
	#分发模型
    Bob_model = model.copy().send('Bob')
    Alice_model = model.copy().send('Alice')

    #分布式训练优化器设置，采用随机梯度下降
    Bob_opt = optim.SGD(params=Bob_model.parameters(),lr=0.1)
    Alice_opt = optim.SGD(params=Alice_model.parameters(),lr=0.1)

    for epoch in range(50):    
        #梯度清0
        Bob_opt.zero_grad()
        Alice_opt.zero_grad()     

        #在Bob上进行训练
        Bob_pred = Bob_model(data_Bob_ptr)
        Bob_loss = ((Bob_pred - targe_Bob_ptr)**2).sum()
        Bob_loss.backward()
        Bob_opt.step()

        #在Alice上面训练
        Alice_pred = Alice_model(data_Alice_ptr)
        Alice_loss = ((Alice_pred - targe_Alice_ptr)**2).sum()
        Alice_loss.backward()
        Alice_opt.step()

        # print(Bob_loss.get().data.item())
        print("第次epoch，avg_loss为：".format(epoch,
                                                (Bob_loss.get().data.item()+Alice_loss.get().data.item())/2))

    with torch.no_grad():
        #回收模型
        Bob_model.get()
        Alice_model.get()
        #进行模型平均后合并
        model.weight.set_((Bob_model.weight.data+Alice_model.weight.data))/2
        model.bias.set_((Bob_model.bias.data+Alice_model.bias.data))/2

	#对Bob的模型进行评估
    pred = Bob_model(torch.tensor([[0,1],[0,1], [1,0], [1,1.]]))
    print(pred.data)

	#对Alice的模型进行评估
    pred = Alice_model(torch.tensor([[0,1],[0,1], [1,0], [1,1.]]))
    print(pred.data)

训练并评估

对整合后的模型进行评估。

train()

pred = model(torch.tensor([[0,1],[0,1], [1,0], [1,1.]]))
print(pred.data)

这里采用一次分发与回收的方案，但是这样我觉得很难整合，如果客户机迭代一次后立刻对模型进行整合，然后再次分发，这算作一个总的迭代过程，是否对模型的整体效果更优？但是这样通信量会增加，训练开销肯定会变大。