import glob import os from pathlib import Path import json import yaml from ultralytics import YOLO import copy import torch # ------------ 新增联邦学习工具函数 ------------ def federated_avg(global_model, client_weights): """联邦平均核心算法""" # 计算总样本数 total_samples = sum(n for _, n in client_weights) if total_samples == 0: raise ValueError("Total number of samples must be positive.") # 获取YOLO底层PyTorch模型参数 global_dict = global_model.model.state_dict() # 提取所有客户端的 state_dict 和对应样本数 state_dicts, sample_counts = zip(*client_weights) for key in global_dict: # 对每一层参数取平均 # if global_dict[key].data.dtype == torch.float32: # global_dict[key].data = torch.stack( # [w[key].float() for w in client_weights], 0 # ).mean(0) # 加权平均 if global_dict[key].dtype == torch.float32: # 只聚合浮点型参数 # 跳过 BatchNorm 层的统计量 if any( x in key for x in ["running_mean", "running_var", "num_batches_tracked"] ): continue # 按照样本数加权求和 weighted_tensors = [ sd[key].float() * (n / total_samples) for sd, n in zip(state_dicts, sample_counts) ] global_dict[key] = torch.stack(weighted_tensors, dim=0).sum(dim=0) # 解决模型参数不匹配问题 # try: # # 加载回YOLO模型 # global_model.model.load_state_dict(global_dict) # except RuntimeError as e: # print('Ignoring "' + str(e) + '"') # 加载回YOLO模型 global_model.model.load_state_dict(global_dict) # 随机选取一个非统计量层进行对比 # sample_key = next(k for k in global_dict if 'running_' not in k) # aggregated_mean = global_dict[sample_key].mean().item() # client_means = [sd[sample_key].float().mean().item() for sd in state_dicts] # print(f"layer: '{sample_key}' Mean after aggregation: {aggregated_mean:.6f}") # print(f"The average value of the layer for each client: {client_means}") # 定义多个关键层 MONITOR_KEYS = [ "model.0.conv.weight", # 输入层卷积 "model.10.conv.weight", # 中间层卷积 "model.22.dfl.conv.weight", # 输出层分类头 ] with open("aggregation_check.txt", "a") as f: f.write("\n=== 参数聚合检查 ===\n") for key in MONITOR_KEYS: if key not in global_dict: continue # 计算聚合后均值 aggregated_mean = global_dict[key].mean().item() # 计算各客户端均值 client_means = [sd[key].float().mean().item() for sd in state_dicts] with open("aggregation_check.txt", "a") as f: f.write(f"层 '{key}' 聚合后均值: {aggregated_mean:.6f}\n") f.write(f"各客户端该层均值差异: {[f'{cm:.6f}' for cm in client_means]}\n") f.write(f"客户端最大差异: {max(client_means) - min(client_means):.6f}\n\n") return global_model # ------------ 修改训练流程 ------------ def federated_train(num_rounds, clients_data): # ========== 新增：初始化指标记录 ========== metrics = { "round": [], "val_mAP": [], # 每轮验证集mAP "train_loss": [], # 每轮平均训练损失 "client_mAPs": [], # 各客户端本地模型在验证集上的mAP "communication_cost": [], # 每轮通信开销（MB） } # 初始化全局模型 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") global_model = YOLO("../yolov8n.yaml").to(device) # 设置类别数 # global_model.model.nc = 1 for _ in range(num_rounds): client_weights = [] client_losses = [] # 记录各客户端的训练损失 # 每个客户端本地训练 for data_path in clients_data: # 统计本地训练样本数 with open(data_path, "r") as f: config = yaml.safe_load(f) # Resolve img_dir relative to the YAML file's location yaml_dir = os.path.dirname(data_path) img_dir = os.path.join( yaml_dir, config.get("train", data_path) ) # 从配置文件中获取图像目录 # print(f"Image directory: {img_dir}") num_samples = (len(glob.glob(os.path.join(img_dir, "*.jpg"))) + len(glob.glob(os.path.join(img_dir, "*.png"))) + len(glob.glob(os.path.join(img_dir, "*.jpeg"))) ) # print(f"Number of images: {num_samples}") # 克隆全局模型 local_model = copy.deepcopy(global_model) # 本地训练（保持你的原有参数设置） results = local_model.train( data=data_path, epochs=4, # 每轮本地训练多少个epoch # save_period=16, imgsz=640, # 图像大小 verbose=False, # 关闭冗余输出 batch=-1, ) # 记录客户端训练损失 # client_loss = results.results_dict['train_loss'] # client_losses.append(client_loss) # 收集模型参数及样本数 client_weights.append( (copy.deepcopy(local_model.model.state_dict()), num_samples) ) # 聚合参数更新全局模型 global_model = federated_avg(global_model, client_weights) # ========== 评估全局模型 ========== # 评估全局模型在验证集上的性能 val_results = global_model.val( data="/mnt/DATA/UAVdataset/data.yaml", # 指定验证集配置文件 imgsz=640, batch=-1, verbose=False, ) val_mAP = val_results.box.map # 获取mAP@0.5 # 计算平均训练损失 # avg_train_loss = sum(client_losses) / len(client_losses) # 计算通信开销（假设传输全部模型参数） model_size = sum(p.numel() * 4 for p in global_model.model.parameters()) / ( 1024**2 ) # MB # 记录到指标容器 metrics["round"].append(_ + 1) metrics["val_mAP"].append(val_mAP) # metrics['train_loss'].append(avg_train_loss) metrics["communication_cost"].append(model_size) # 打印当前轮次结果 with open("aggregation_check.txt", "a") as f: f.write(f"\n[Round {_ + 1}/{num_rounds}]") f.write(f"Validation mAP@0.5: {val_mAP:.4f}") # f.write(f"Average Train Loss: {avg_train_loss:.4f}") f.write(f"Communication Cost: {model_size:.2f} MB\n") return global_model, metrics # ------------ 使用示例 ------------ if __name__ == "__main__": # 联邦训练配置 clients_config = [ "/mnt/DATA/uav_dataset_fed/train1/train1.yaml", # 客户端1数据路径 "/mnt/DATA/uav_dataset_fed/train2/train2.yaml", # 客户端2数据路径 ] # 运行联邦训练 final_model, metrics = federated_train(num_rounds=40, clients_data=clients_config) # 保存最终模型 final_model.save("yolov8n_federated.pt") # final_model.export(format="onnx") # 导出为ONNX格式 with open("metrics.json", "w") as f: json.dump(metrics, f, indent=4)