Graduation-Project/federated_learning/yolov8_fed.py

import glob
import os
from pathlib import Path
import json
from pydoc import cli
from threading import local

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.")

    # DEBUG: global_dict
    # print(global_model)

    # 获取YOLO底层PyTorch模型参数
    global_dict = global_model.model.state_dict()
    # 提取所有客户端的 state_dict 和对应样本数
    state_dicts, sample_counts = zip(*client_weights)

    # 克隆参数并脱离计算图
    global_dict_copy = {
        k: v.clone().detach().requires_grad_(False) for k, v in global_dict.items()
    }

    # 聚合可训练且存在的参数
    for key in global_dict_copy:
        # if global_dict_copy[key].dtype != torch.float32:
        #     continue
        # if any(
        #     x in key for x in ["running_mean", "running_var", "num_batches_tracked"]
        # ):
        #     continue
        # 检查所有客户端是否包含当前键
        all_clients_have_key = all(key in sd for sd in state_dicts)
        if all_clients_have_key:
            # 计算每个客户端的加权张量
            # weighted_tensors = [
            #     client_state[key].float() * (sample_count / total_samples)
            #     for client_state, sample_count in zip(state_dicts, sample_counts)
            # ]
            weighted_tensors = []
            for client_state, sample_count in zip(state_dicts, sample_counts):
                weight = sample_count / total_samples  # 计算权重
                weighted_tensor = client_state[key].float() * weight  # 加权张量
                weighted_tensors.append(weighted_tensor)
            # 聚合加权张量并更新全局参数
            global_dict_copy[key] = torch.stack(weighted_tensors, dim=0).sum(dim=0)

        # else:
        #     print(f"错误: 键 {key} 在部分客户端缺失，已保留全局参数")
        # 终止训练或记录日志
        # raise KeyError(f"键 {key} 缺失")

    # 加载回YOLO模型
    global_model.model.load_state_dict(global_dict_copy, strict=True)

    # global_model.model.train()
    # with torch.no_grad():
    #     global_model.model.load_state_dict(global_dict_copy, strict=True)

    # 定义多个关键层
    MONITOR_KEYS = [
        "model.0.conv.weight",
        "model.1.conv.weight",
        "model.3.conv.weight",
        "model.5.conv.weight",
        "model.7.conv.weight",
        "model.9.cv1.conv.weight",
        "model.12.cv1.conv.weight",
        "model.15.cv1.conv.weight",
        "model.18.cv1.conv.weight",
        "model.21.cv1.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
        # if not all(key in sd for sd in state_dicts):
        #     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("/home/image1325/DATA/Graduation-Project/federated_learning/yolov8n.yaml")
        .load("/home/image1325/DATA/Graduation-Project/federated_learning/yolov8n.pt")
        .to(device)
    )
    global_model.model.model[-1].nc = 1  # 设置检测类别数为1
    # global_model.model.train.ema.enabled = False

    # 克隆全局模型
    local_model = copy.deepcopy(global_model)

    for _ in range(num_rounds):
        client_weights = []
        # 各客户端的训练损失
        # client_losses = []

        # DEBUG: 检查全局模型参数
        # global_dict = global_model.model.state_dict()
        # print(global_dict.keys())

        # 每个客户端本地训练
        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.model.load_state_dict(
                global_model.model.state_dict(), strict=True
            )

            # 本地训练（保持你的原有参数设置）
            local_model.train(
                name=f"train{_ + 1}",  # 当前轮次
                data=data_path,
                # model=local_model,
                epochs=16,  # 每轮本地训练多少个epoch
                # save_period=16,
                imgsz=768,  # 图像大小
                verbose=False,  # 关闭冗余输出
                batch=-1,  # 批大小
                workers=6,  # 工作线程数
            )

            # 记录客户端训练损失
            # client_loss = results.results_dict['train_loss']
            # client_losses.append(client_loss)

            # 收集模型参数及样本数
            client_weights.append((local_model.model.state_dict(), num_samples))

        # 聚合参数更新全局模型
        global_model = federated_avg(global_model, client_weights)

        # DEBUG: 检查全局模型参数
        # keys = global_model.model.state_dict().keys()

        # ========== 评估全局模型 ==========
        # 复制全局模型以避免在评估时修改参数
        val_model = copy.deepcopy(global_model)
        # 评估全局模型在验证集上的性能
        with torch.no_grad():
            val_results = val_model.val(
                data="/mnt/DATA/uav_dataset_old/UAVdataset/fed_data.yaml",  # 指定验证集配置文件
                imgsz=768,  # 图像大小
                batch=16,  # 批大小
                verbose=False,  # 关闭冗余输出
            )
        # 丢弃评估模型
        del val_model

        # DEBUG: 检查全局模型参数
        # if keys != global_model.model.state_dict().keys():
        #     print("模型参数不一致！")

        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}]\n")
            f.write(f"Validation mAP@0.5: {val_mAP:.4f}\n")
            # f.write(f"Average Train Loss: {avg_train_loss:.4f}")
            f.write(f"Communication Cost: {model_size:.2f} MB\n\n")

    return global_model, metrics


if __name__ == "__main__":
    # 联邦训练配置
    clients_config = [
        "/mnt/DATA/uav_fed/train1/train1.yaml",  # 客户端1数据路径
        "/mnt/DATA/uav_fed/train2/train2.yaml",  # 客户端2数据路径
    ]

    # 使用本地数据集进行测试
    # clients_config = [
    #     "/home/image1325/DATA/Graduation-Project/dataset/train1/train1.yaml",
    #     "/home/image1325/DATA/Graduation-Project/dataset/train2/train2.yaml",
    # ]

    # 运行联邦训练
    final_model, metrics = federated_train(num_rounds=10, 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)