def create_balanced_tree():
"""Tree: {None: 1}{1: 2}{1: 3}{2:4}{3: 6}{3: 7}"""
tree = Tree()
tree.insert(Node(4, parent_name=None))
tree.insert(Node(2, parent_name=4))
tree.insert(Node(6, parent_name=4))
tree.insert(Node(3, parent_name=2))
tree.insert(Node(5, parent_name=6))
tree.insert(Node(7, parent_name=6))
return tree
def create_complete_tree():
"""Tree: {None: 4}{4: 2}{4: 6}{2: 1}{2: 3}{6: 5}"""
tree = Tree()
tree.insert(Node(4, parent_name=None))
tree.insert(Node(2, parent_name=4))
tree.insert(Node(6, parent_name=4))
tree.insert(Node(1, parent_name=2))
tree.insert(Node(3, parent_name=2))
tree.insert(Node(5, parent_name=6))
return tree
def test_tree_print_with_two_children(self):
root = Node(name="first_level")
t = Tree(root=root, name="root")
left = Node(name="second_level_one")
right = Node(name="second_level_two")
t.root.add_child(left)
t.root.add_child(right)
expected = (
"<Tree 'root'>\n<Node first_level>\n\t<Node second_level_one>\n\t")
expected += "<Node second_level_two>" # ...
self.assertEqual(str(t), expected)
def test_tree_constructor(self):
t = Tree(name="spruce")
self.assertEqual(t.name, "spruce")
self.assertIsNotNone(t.root)
def test_tree_print_with_one_child(self):
t = Tree()
left = Node(name="left")
t.root.add_child(left)
expected = "<Tree 'None'>\n<Node None>\n\t<Node left>"
self.assertEqual(str(t), expected)
def test_tree_print_just_root(self):
t = Tree()
expected = "<Tree 'None'>\n<Node None>"
self.assertEqual(str(t), expected)
def ScoreAverage(dps):
"""非树状聚合方式（即一起聚合） 计算SV方式：平均outputs在测试集上精度"""
shapleyValue = ShapleyValue()
tp = ThirdParty()
db = DataBuyer(get_net())
'''dps = []
for i in range(params.provider_num):
net = get_net()
dataloader = get_data_loader(i)
dps.append(DataProvider(net, dataloader))
print('读取模型完成')'''
# 随机聚合顺序
'''order_rand = random_order(params.provider_num)
print('聚合顺序', order_rand)
'''
# 构成树节点 放入tree_list
tree_list = []
for i in range(params.provider_num):
tree_list.append(Tree(i, dps[i]))
# 先在本地数据集上训练至收敛----------------
# '''
for i in range(params.provider_num):
print("客户端", i, "预训练")
for j in range(params.local_time):
tree_list[i].provider.train()
# '''
# 计算SV-------------------
print('开始计算SV')
shapleyValue.v_way = 'score_avg' # 计算v的方式fedavg和score_avg
SVs = shapleyValue.cal_SV_all(tree_list)
print("算得各个SV值：")
for i in range(params.provider_num):
tree_list[i].sv = SVs[i]
print(SVs[i])
# 找出SV>0的聚合-----------------
positive_list = []
for i in range(params.provider_num):
if SVs[i] > 0:
print(i, "SV>0并加入")
positive_list.append(tree_list[i])
net, acc = fedavg(positive_list, 100)
print("聚合后精度", acc)
# 写入txt
txt_dir = params.dataset_division_testno + '/22.txt'
write_txt(tree_list, 0, acc, txt_dir)
# 把v_all写入txt
v_all = shapleyValue.v_all
print(v_all)
npy_dir = params.dataset_division_testno + '/ScoreAverage_v_all_2.npy'
write_npy_v_all(v_all, npy_dir)
def Original(dps):
"""原本聚合方式：FedAvg +计算SV"""
shapleyValue = ShapleyValue()
db = DataBuyer(get_net())
'''dps = []
for i in range(params.provider_num):
net = get_net()
dataloader = get_data_loader(i)
dps.append(DataProvider(net, dataloader))'''
# 构成树节点 放入tree_list
tree_list = []
for i in range(params.provider_num):
tree_list.append(Tree(i, dps[i]))
# 预训练
# _, p_fed = fedavg(tree_list)
num_node = len(tree_list)
# 先在本地数据集上训练至收敛----------------
# '''
for i in range(params.provider_num):
print("客户端", i, "预训练")
for j in range(params.local_time):
tree_list[i].provider.train()
# '''
#
print('开始计算SV')
shapleyValue.v_way = 'fedavg' # 计算v的方式fedavg和score_avg
SVs = shapleyValue.cal_SV_all(tree_list)
# 所有聚合后pab
v_all = shapleyValue.v_all
for i in range(num_node):
tree_list[i].sv = SVs[i]
print(SVs[i])
# 最后剩一个节点为根
root = tree_list[0]
root.B = db.B
# 根据树分配B
all_B(root)
# 根节点精确度
p_root = shapleyValue.root_p
# SV写入txt
txt_dir = params.dataset_division_testno + '/21.txt'
write_txt(tree_list, 0, p_root, txt_dir)
# 把v_all写入txt
v_all = shapleyValue.v_all
print(v_all)
npy_dir = params.dataset_division_testno + '/Original_v_all_2.npy'
write_npy_v_all(v_all, npy_dir)
# 第三方解密，发送结果给DP、DB
return tree_list
def CollaborativeModelling(_tree_list=None):
"""树状聚合方式 +计算SV"""
shapleyValue = ShapleyValue()
tp = ThirdParty()
db = DataBuyer(get_net())
dps = []
for i in range(params.provider_num):
net = get_net()
dataloader = get_data_loader(i)
dps.append(DataProvider(net, dataloader))
print('读取模型完成')
# 随机聚合顺序
'''order_rand = random_order(params.provider_num)
print('聚合顺序', order_rand)
'''
# 构成树节点 放入tree_list
tree_list = []
if _tree_list is not None:
for i in range(params.provider_num):
tree_list.append(_tree_list[i])
else:
for i in range(params.provider_num):
tree_list.append(Tree(i, dps[i]))
"""# 第三方生成密匙，传给DP、DB
public_key, private_key = tp.generate_key()
# DP加密model，发给DB
for i in range(params.provider_num):
dps[i].enctypt()
# 聚合前先FedAvg p_fed为fedavg的精度
_, p_fed = fedavg(tree_list)
"""
# 开始多次FedAvg、聚合
last_node = tree_list[0] # 上一次最优节点
next_node_no = 1 # 接下来要聚合的开始节点编号
node_K_list = [last_node]
while next_node_no < params.provider_num:
# print('len(node_K_list[0].children)', len(node_K_list[0].children))
# 要聚合的K个节点
num = 0
while num < params.K - 1 and next_node_no < params.provider_num:
node_K_list.append(tree_list[next_node_no])
next_node_no += 1
num += 1
# K个provider聚合 可能不足K个
num_node = len(node_K_list)
# DB计算特征函数v，发送给第三方
print('开始计算SV')
SVs = shapleyValue.cal_SV_all(node_K_list)
for i in range(num_node):
node_K_list[i].sv = SVs[i]
print(SVs[i])
# 判断是否聚合
num_aggregation = 0
for i in range(num_node):
if node_K_list[i].if_aggregation():
num_aggregation += 1
if num_aggregation == num_node: # 全部同意聚合
# 用聚合的模型建树
net = shapleyValue.root_net
# 暂时用第一个孩子的dataloader做聚合节点的dataloader
p = DataProvider(net,
dataloader=get_data_loader(node_K_list[0].p_no))
node = Tree(node_K_list[0].p_no, p)
for i in range(num_node):
node.children.append(node_K_list[i])
# 记录上一次聚合的节点
node_K_list = [node]
else:
# 选出SV最大的节点做根
max_node = node_K_list[0]
max_sv = node_K_list[0].sv
for i in range(1, num_node):
if node_K_list[i].sv > max_sv:
max_node = node_K_list[i]
max_sv = node_K_list[i].sv
# 记录上一次最优的节点
node_K_list = [max_node]
# 最后剩一个节点为根
root = node_K_list[0]
root.B = db.B
# 根据树分配B
all_B(root)
# 根节点精确度
p_root = shapleyValue.root_p
# 写入txt
txt_dir = params.dataset_division_testno + '/10.txt'
write_txt(tree_list, 0, p_root, txt_dir)
def create_unbalanced_tree():
"""Tree: {None: 1}{1: 2}{1: 3}{3: 6}{3: 7}{6: 8}"""
tree = Tree()
log.debug(tree.get_root())
tree.insert(Node(4, parent_name=None))
tree.insert(Node(2, parent_name=4))
tree.insert(Node(6, parent_name=4))
tree.insert(Node(1, parent_name=2))
tree.insert(Node(5, parent_name=6))
tree.insert(Node(7, parent_name=6))
tree.insert(Node(8, parent_name=7))
tree.insert(Node(9, parent_name=8))
return tree