Anchor-x/Graduation-Project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: Anchor-x:2e7cf69512d7e5717758c493f82e370b4881fa09
Branches Tags
Anchor-x:main
..
compare: Anchor-x:f127ae2852a15fad18cfbcb0fc826715e41184e1
Branches Tags
Anchor-x:main

No commits in common. "2e7cf69512d7e5717758c493f82e370b4881fa09" and "f127ae2852a15fad18cfbcb0fc826715e41184e1" have entirely different histories.
2e7cf69512 ... f127ae2852

4 changed files with 71 additions and 149 deletions
Show Stats Expand all files Collapse all files

31
README.md
View File

@ -2,34 +2,7 @@
毕业设计：基于YOLO和图像融合技术的无人机检测系统及安全性研究 毕业设计：基于YOLO和图像融合技术的无人机检测系统及安全性研究
Linux 运行联邦训练 Linux 运行训练
```bash ```bash
cd federated_learning nohup python -u yolov8_fed.py >> runtime.log 2>&1 &
```
```bash
nohup python -u yolov8_fed.py > runtime.log 2>&1 &
```
Linux 运行集中训练
```bash
cd yolov8
```
```bash
nohup python -u yolov8_train.py > runtime.log 2>&1 &
```
实时监控日志文件
```bash
tail -f runtime.log
```
运行图像融合配准代码
```bash
cd image_fusion
```
```bash
python Image_Registration_test.py
``` ```

189
federated_learning/yolov8_fed.py
View File

@ -2,8 +2,6 @@ import glob
import os import os
from pathlib import Path from pathlib import Path
import json import json
from pydoc import cli
from threading import local
import yaml import yaml
from ultralytics import YOLO from ultralytics import YOLO
@ -19,83 +17,66 @@ def federated_avg(global_model, client_weights):
if total_samples == 0: if total_samples == 0:
raise ValueError("Total number of samples must be positive.") raise ValueError("Total number of samples must be positive.")
# DEBUG: global_dict
# print(global_model)
# 获取YOLO底层PyTorch模型参数 # 获取YOLO底层PyTorch模型参数
global_dict = global_model.model.state_dict() global_dict = global_model.model.state_dict()
# 提取所有客户端的 state_dict 和对应样本数 # 提取所有客户端的 state_dict 和对应样本数
state_dicts, sample_counts = zip(*client_weights) state_dicts, sample_counts = zip(*client_weights)
# 克隆参数并脱离计算图 for key in global_dict:
global_dict_copy = { # 对每一层参数取平均
k: v.clone().detach().requires_grad_(False) for k, v in global_dict.items() # 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)
# 聚合可训练且存在的参数 # 加权平均
for key in global_dict_copy: if global_dict[key].dtype == torch.float32: # 只聚合浮点型参数
# if global_dict_copy[key].dtype != torch.float32: # 跳过 BatchNorm 层的统计量
# continue if any(
# if any( x in key for x in ["running_mean", "running_var", "num_batches_tracked"]
# x in key for x in ["running_mean", "running_var", "num_batches_tracked"] ):
# ): continue
# continue # 按照样本数加权求和
# 检查所有客户端是否包含当前键 weighted_tensors = [
all_clients_have_key = all(key in sd for sd in state_dicts) sd[key].float() * (n / total_samples)
if all_clients_have_key: for sd, n in zip(state_dicts, sample_counts)
# 计算每个客户端的加权张量 ]
# weighted_tensors = [ global_dict[key] = torch.stack(weighted_tensors, dim=0).sum(dim=0)
# 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} 在部分客户端缺失，已保留全局参数") # try:
# 终止训练或记录日志 # # 加载回YOLO模型
# raise KeyError(f"键 {key} 缺失") # global_model.model.load_state_dict(global_dict)
# except RuntimeError as e:
# print('Ignoring "' + str(e) + '"')
# 加载回YOLO模型 # 加载回YOLO模型
global_model.model.load_state_dict(global_dict_copy, strict=True) global_model.model.load_state_dict(global_dict)
# global_model.model.train() # 随机选取一个非统计量层进行对比
# with torch.no_grad(): # sample_key = next(k for k in global_dict if 'running_' not in k)
# global_model.model.load_state_dict(global_dict_copy, strict=True) # 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 = [ MONITOR_KEYS = [
"model.0.conv.weight", "model.0.conv.weight", # 输入层卷积
"model.1.conv.weight", "model.10.conv.weight", # 中间层卷积
"model.3.conv.weight", "model.22.dfl.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: with open("aggregation_check.txt", "a") as f:
f.write("\n=== 参数聚合检查 ===\n") f.write("\n=== 参数聚合检查 ===\n")
for key in MONITOR_KEYS: for key in MONITOR_KEYS:
# if key not in global_dict: if key not in global_dict:
# continue continue
# if not all(key in sd for sd in state_dicts):
# continue
# 计算聚合后均值 # 计算聚合后均值
aggregated_mean = global_dict[key].mean().item() aggregated_mean = global_dict[key].mean().item()
# 计算各客户端均值 # 计算各客户端均值
client_means = [sd[key].float().mean().item() for sd in state_dicts] client_means = [sd[key].float().mean().item() for sd in state_dicts]
with open("aggregation_check.txt", "a") as f: with open("aggregation_check.txt", "a") as f:
f.write(f"层 '{key}' 聚合后均值: {aggregated_mean:.6f}\n") f.write(f"层 '{key}' 聚合后均值: {aggregated_mean:.6f}\n")
f.write(f"各客户端该层均值差异: {[f'{cm:.6f}' for cm in client_means]}\n") f.write(f"各客户端该层均值差异: {[f'{cm:.6f}' for cm in client_means]}\n")
@ -106,35 +87,24 @@ def federated_avg(global_model, client_weights):
# ------------ 修改训练流程 ------------ # ------------ 修改训练流程 ------------
def federated_train(num_rounds, clients_data): def federated_train(num_rounds, clients_data):
# ========== 初始化指标记录 ========== # ========== 新增：初始化指标记录 ==========
metrics = { metrics = {
"round": [], "round": [],
"val_mAP": [], # 每轮验证集mAP "val_mAP": [], # 每轮验证集mAP
# "train_loss": [], # 每轮平均训练损失 "train_loss": [], # 每轮平均训练损失
"client_mAPs": [], # 各客户端本地模型在验证集上的mAP "client_mAPs": [], # 各客户端本地模型在验证集上的mAP
"communication_cost": [], # 每轮通信开销（MB） "communication_cost": [], # 每轮通信开销（MB）
} }
# 初始化全局模型 # 初始化全局模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
global_model = ( global_model = YOLO("../yolov8n.yaml").to(device)
YOLO("/home/image1325/DATA/Graduation-Project/federated_learning/yolov8n.yaml") # 设置类别数
.load("/home/image1325/DATA/Graduation-Project/federated_learning/yolov8n.pt") # global_model.model.nc = 1
.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): for _ in range(num_rounds):
client_weights = [] client_weights = []
# 各客户端的训练损失 client_losses = [] # 记录各客户端的训练损失
# client_losses = []
# DEBUG: 检查全局模型参数
# global_dict = global_model.model.state_dict()
# print(global_dict.keys())
# 每个客户端本地训练 # 每个客户端本地训练
for data_path in clients_data: for data_path in clients_data:
@ -148,28 +118,23 @@ def federated_train(num_rounds, clients_data):
) # 从配置文件中获取图像目录 ) # 从配置文件中获取图像目录
# print(f"Image directory: {img_dir}") # print(f"Image directory: {img_dir}")
num_samples = ( num_samples = (len(glob.glob(os.path.join(img_dir, "*.jpg")))
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, "*.png")))
+ len(glob.glob(os.path.join(img_dir, "*.jpeg"))) + len(glob.glob(os.path.join(img_dir, "*.jpeg")))
) )
# print(f"Number of images: {num_samples}") # print(f"Number of images: {num_samples}")
local_model.model.load_state_dict( # 克隆全局模型
global_model.model.state_dict(), strict=True local_model = copy.deepcopy(global_model)
)
# 本地训练（保持你的原有参数设置） # 本地训练（保持你的原有参数设置）
local_model.train( results = local_model.train(
name=f"train{_ + 1}", # 当前轮次
data=data_path, data=data_path,
# model=local_model, epochs=4, # 每轮本地训练多少个epoch
epochs=16, # 每轮本地训练多少个epoch
# save_period=16, # save_period=16,
imgsz=768, # 图像大小 imgsz=640, # 图像大小
verbose=False, # 关闭冗余输出 verbose=False, # 关闭冗余输出
batch=-1, # 批大小 batch=-1,
workers=6, # 工作线程数
) )
# 记录客户端训练损失 # 记录客户端训练损失
@ -177,32 +142,21 @@ def federated_train(num_rounds, clients_data):
# client_losses.append(client_loss) # client_losses.append(client_loss)
# 收集模型参数及样本数 # 收集模型参数及样本数
client_weights.append((local_model.model.state_dict(), num_samples)) client_weights.append(
(copy.deepcopy(local_model.model.state_dict()), num_samples)
)
# 聚合参数更新全局模型 # 聚合参数更新全局模型
global_model = federated_avg(global_model, client_weights) 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 = global_model.val(
val_results = val_model.val( data="/mnt/DATA/UAVdataset/data.yaml", # 指定验证集配置文件
data="/mnt/DATA/uav_dataset_old/UAVdataset/fed_data.yaml", # 指定验证集配置文件 imgsz=640,
imgsz=768, # 图像大小 batch=-1,
batch=16, # 批大小 verbose=False,
verbose=False, # 关闭冗余输出 )
)
# 丢弃评估模型
del val_model
# DEBUG: 检查全局模型参数
# if keys != global_model.model.state_dict().keys():
# print("模型参数不一致！")
val_mAP = val_results.box.map # 获取mAP@0.5 val_mAP = val_results.box.map # 获取mAP@0.5
# 计算平均训练损失 # 计算平均训练损失
@ -220,29 +174,24 @@ def federated_train(num_rounds, clients_data):
metrics["communication_cost"].append(model_size) metrics["communication_cost"].append(model_size)
# 打印当前轮次结果 # 打印当前轮次结果
with open("aggregation_check.txt", "a") as f: with open("aggregation_check.txt", "a") as f:
f.write(f"\n[Round {_ + 1}/{num_rounds}]\n") f.write(f"\n[Round {_ + 1}/{num_rounds}]")
f.write(f"Validation mAP@0.5: {val_mAP:.4f}\n") f.write(f"Validation mAP@0.5: {val_mAP:.4f}")
# f.write(f"Average Train Loss: {avg_train_loss:.4f}") # f.write(f"Average Train Loss: {avg_train_loss:.4f}")
f.write(f"Communication Cost: {model_size:.2f} MB\n\n") f.write(f"Communication Cost: {model_size:.2f} MB\n")
return global_model, metrics return global_model, metrics
# ------------ 使用示例 ------------
if __name__ == "__main__": if __name__ == "__main__":
# 联邦训练配置 # 联邦训练配置
clients_config = [ clients_config = [
"/mnt/DATA/uav_fed/train1/train1.yaml", # 客户端1数据路径 "/mnt/DATA/uav_dataset_fed/train1/train1.yaml", # 客户端1数据路径
"/mnt/DATA/uav_fed/train2/train2.yaml", # 客户端2数据路径 "/mnt/DATA/uav_dataset_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, metrics = federated_train(num_rounds=40, clients_data=clients_config)
# 保存最终模型 # 保存最终模型
final_model.save("yolov8n_federated.pt") final_model.save("yolov8n_federated.pt")

0
federated_learning/yolov8.yaml → yolov8.yaml
View File

BIN
yolov8n.pt Normal file
View File

Binary file not shown.