Pysyft实验1

# 联邦学习案例学习
# 通过模型平均的联邦学习案例特点：
# 1.模型平均使用多个模型的权重平均
# 2.通过可信的第三者聚合权重
# 3.训练使用线性模型进行分类
#
# 操作步骤
# 1.创建两个数据的owner(bob和alice)
# 2.创建模型model
# 3.发送模型的copy版本到alice和bob
# 4.并行训练alice和bob的模型
# 5.发送更新的模型给secure worker
# 6.加权平均模型
import torch
import syft as sy
import copy
hook = sy.TorchHook(torch)
from torch import nn, optim

# create a couple workers/虚拟工人
bob = sy.VirtualWorker(hook, id="bob")
alice = sy.VirtualWorker(hook, id="alice")
secure_worker = sy.VirtualWorker(hook, id="secure_worker")


# A Toy Dataset/简单的数据集
data = torch.tensor([[0,0],[0,1],[1,0],[1,1.]], requires_grad=True)
target = torch.tensor([[0],[0],[1],[1.]], requires_grad=True)

# get pointers to training data on each worker by
# 这句话的意思是获取对每个工作节点上训练数据的指针。
# 在联邦学习中，由于数据通常存储在不同的地方（比如不同的设备或服务器），为了训练模型，
# 我们需要在每个位置获取对数据的引用或指针，而不是将数据直接传输到中心服务器。
# sending some training data to bob and alice
bobs_data = data[0:2].send(bob)
bobs_target = target[0:2].send(bob)

alices_data = data[2:].send(alice)
alices_target = target[2:].send(alice)

# Iniitalize A Toy Model模型初始化
model = nn.Linear(2,1)#这行代码初始化了一个简单的线性模型，使用PyTorch的nn.Linear类，指定输入特征的维度为2，输出特征的维度为1
iter_num = 20#这行代码初始化了一个名为iter_num的变量，值为20，用于表示迭代的次数。
parallel_worker_iter_num=10#这行代码初始化了一个名为parallel_worker_iter_num的变量，值为10，用于表示并行工作者的迭代次数。


#开始迭代训练模型
for j in range(iter_num):
    bobs_model = model.copy().send(bob)#复制模型并发给bob
    alices_model = model.copy().send(alice)#复制模型并发送给 alice

    #创建优化器
    bobs_opt = optim.SGD(params=bobs_model.parameters(), lr=0.1)
    alices_opt = optim.SGD(params=alices_model.parameters(), lr=0.1)

    # 并行训练 bob 和 alice 的模型
    for i in range(parallel_worker_iter_num):
        # Train Bob's Model
        bobs_opt.zero_grad()# 梯度清零
        bobs_pred = bobs_model(bobs_data)# 预测
        bobs_loss = ((bobs_pred - bobs_target) ** 2).sum()# 计算损失
        bobs_loss.backward()# 反向传播

        bobs_opt.step()# 更新模型参数
        bobs_loss = bobs_loss.get().data# 获取损失值

        # Train Alice's Model
        alices_opt.zero_grad()# 梯度清零
        alices_pred = alices_model(alices_data)# 预测
        alices_loss = ((alices_pred - alices_target) ** 2).sum()# 计算损失
        alices_loss.backward()# 反向传播

        alices_opt.step()# 更新模型参数
        alices_loss = alices_loss.get().data# 获取损失值

    # 将 Alice 和 Bob 的模型参数移动到安全的虚拟工人 secure_worker
    alices_model.move(secure_worker)
    bobs_model.move(secure_worker)
    # 使用加权平均更新主模型的参数
    with torch.no_grad():
        # 更新主模型权重
        model.weight.set_(((alices_model.weight.data + bobs_model.weight.data) / 2).get())
        # 更新主模型偏置
        model.bias.set_(((alices_model.bias.data + bobs_model.bias.data) / 2).get())
    print(j,"iter:","Bob_loss:" + str(bobs_loss) + " Alice_loss:" + str(alices_loss))
# 使用整体模型进行预测
preds = model(data)# 进行预测
loss = ((preds - target) ** 2).sum()# 计算总损失值

print("preds:",preds) # 输出预测结果
print("target:",target)# 输出目标值
print("loss:",loss.data)# 输出损失值

代码参考链接：pysyft实战笔记 - 知乎 (zhihu.com)

运行结果：