| 1234567891011121314151617181920212223242526272829303132333435 |
- from peft import (
- set_peft_model_state_dict,
- )
- import torch
- import os
- from torch.nn.functional import normalize
- # 联邦平均算法
- def FedAvg(model, selected_clients_set, output_dir, local_dataset_len_dict, epoch):
- # 对各个客户端的本地数据集大小进行归一化，作为权重的基础
- # 这里将每个客户端的数据集大小转换为张量，并按照第0维进行 L1 归一化
- weights_array = normalize(
- torch.tensor([local_dataset_len_dict[client_id] for client_id in selected_clients_set],
- dtype=torch.float32),
- p=1, dim=0)
- # 遍历选定的客户端集合
- for k, client_id in enumerate(selected_clients_set):
- # 对每个选定的客户端，加载其训练得到的模型权重
- # 构造每个客户端权重的文件路径，并加载权重
- single_output_dir = os.path.join(output_dir, str(epoch), "local_output_{}".format(client_id),
- "pytorch_model.bin")
- single_weights = torch.load(single_output_dir)
- # 如果是第一个客户端，则将其权重乘以对应的归一化权重
- # 否则，将其权重乘以归一化权重后累加到总权重上
- if k == 0:
- weighted_single_weights = {key: single_weights[key] * (weights_array[k]) for key in
- single_weights.keys()}
- else:
- weighted_single_weights = {key: weighted_single_weights[key] + single_weights[key] * (weights_array[k])
- for key in
- single_weights.keys()}
- # 使用计算得到的加权平均权重，设置全局模型的状态
- set_peft_model_state_dict(model, weighted_single_weights, "default")
- return model
|
