# YOLO一体化平台搭建：训练、测试、优化、部署全栈指南

# YOLO一体化平台搭建：训练、测试、优化、部署全栈指南

从零构建一个集训练、测试、剪枝、导出、可视化于一体的YOLO智能平台，涵盖PyQt5 GUI开发、数据库管理、TensorBoard集成与模型优化全流程。

第1章 引言：为什么要搭建一体化平台？

1.1 当前YOLO开发流程的痛点

在深度学习目标检测的日常开发中，我们常常面临以下困境：

• 工具链割裂：训练用命令行，测试写脚本，剪枝换一套工具，导出又得查文档——每次切换上下文都消耗大量精力。
• 实验管理混乱：不同超参数组合的训练结果散落在各个文件夹，时间一长连自己都分不清哪个配置对应哪个结果。
• 缺乏可视化反馈：训练过程只能看终端滚动的loss数字，无法直观感受模型学习动态。
• 模型优化门槛高：剪枝、量化、导出不同格式需要记忆大量命令行参数，稍有不慎就会踩坑。
• 历史追溯困难：想对比两次训练的效果，只能手动翻日志、画表格，效率极低。

这些问题在大规模实验和团队协作中会被进一步放大。

1.2 一体化平台的设计目标

针对上述痛点，我们设计并实现了一个“YOLO一体化智能训练系统”，核心目标包括：

• 全流程覆盖：在一个GUI界面中完成数据准备、模型训练、指标监控、效果测试、模型剪枝、格式导出。
• 超参数可视化调节**：所有训练超参数（学习率、数据增强、损失权重等）均通过控件调节，无需记忆命令行参数。
• 自动实验管理：训练过程自动记录到SQLite数据库，支持历史查询与指标曲线绘制。
• 实时监控反馈：集成TensorBoard，训练指标实时刷新，同时提供简洁的文本指标显示。
• 模型优化一体化：内置结构化剪枝（基于torch_pruning）和多格式导出（ONNX、OpenVINO）。
• 智能路径解析：自动从data.yaml中解析训练图片目录，支持手动覆写，测试时随机抽取图片验证模型效果。

1.3 系统技术栈概览

1.4 本文结构

我们将从架构设计、核心模块实现、关键技术难点、优化与部署、实战案例五个维度，深入剖析该平台的构建过程。文章末尾附带完整代码注释，适合有一定Python和深度学习基础的开发者阅读。

第2章 系统架构与模块设计

2.1 整体架构分层

系统采用经典的四层架构，确保各模块职责清晰、耦合度低：

2.2 核心模块划分

系统包含以下核心模块：

• TrainWorker（训练工作线程）：继承QThread，负责异步执行YOLO训练任务，通过信号与主线程通信。
• MetricsDatabase（指标数据库）：管理训练会话和epoch指标的SQLite操作。
• HistoryWidget（历史记录组件）：查询数据库并绘制Matplotlib曲线。
• ModelManagerWidget（模型管理组件：执行剪枝和模型导出操作。
• PruneWorker（剪枝工作线程）：异步执行结构化剪枝。
• YOLOTrainGUI（主窗口）：整合所有子模块，负责用户交互和整体流程控制。

2.3 数据流设计

系统核心数据流如下：
用户配置参数 → 训练启动
    ↓
TrainWorker初始化 → 创建数据库会话
    ↓
训练循环（每个epoch回调）
    ├── 指标提取 → 数据库存储
    ├── 指标提取 → TensorBoard日志
    ├── 指标提取 → CSV文件
    └── 指标提取 → 实时UI更新
    ↓
训练完成 → 更新会话结束时间 → 触发测试

这种设计确保了训练过程的每个关键节点都有数据持久化和可视化反馈，方便后续追溯和分析。

第3章 环境搭建与依赖配置

3.1 基础环境要求

• 操作系统：Windows 10/11、Ubuntu 20.04+、macOS 12+
• Python版本：3.8 ~ 3.11（推荐3.10）
• CUDA：11.8+（如需GPU训练）
• 显存：建议8GB以上（训练YOLOv26-nano约需4GB）

3.2 核心依赖安装

基础深度学习环境配置如下所示：

# 基础深度学习环境
pip install torch torchvision torchaudio –index-url https://download.pytorch.org/whl/cu118
pip install ultralytics==8.2.0

# GUI相关
pip install PyQt5==5.15.9 PyQtWebEngine==5.15.6

# 数据处理与可视化
pip install opencv-python==4.9.0.80
pip install numpy==1.24.3
pip install matplotlib==3.8.0
pip install tensorboard==2.16.0

# 模型优化
pip install torch-pruning==1.1.0
pip install thop==0.1.1.post2209072238

# 数据库
pip install sqlite3 # Python内置，无需额外安装

# 导出工具（可选）
pip install onnx==1.15.0
pip install onnxruntime==1.17.1
pip install openvino==2024.0.0

3.3 验证安装

import torch
import ultralytics
import PyQt5
import cv2
import matplotlib
import tensorboard
import torch_pruning

print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
print(f"Ultralytics版本: {ultralytics.__version__}")
print(f"OpenCV版本: {cv2.__version__}")

3.4 数据集准备

系统需要标准YOLO格式的数据集，目录结构如下：

dataset/
├── data.yaml          # 数据集配置文件
├── images/
│   ├── train/         # 训练图片
│   └── val/           # 验证图片
└── labels/
    ├── train/         # 训练标签
    └── val/           # 验证标签

data.yaml示例：

path: ./dataset
train: images/train
val: images/val

nc: 80
names: ['person', 'bicycle', 'car', …] # COCO类别名

第4章 数据库模块：训练指标持久化

4.1 为什么需要数据库？

在深度学习实验管理中，我们通常需要回答以下问题：

哪次训练效果最好？对应的超参数是什么？
不同学习率下loss收敛曲线有何差异？
某个epoch的详细指标数值是多少？

传统做法是保存CSV或JSON文件，但存在以下问题：

文件散落，难以统一查询
缺乏结构化查询能力（如按时间排序、按数据集筛选）
跨实验对比需要手动复制粘贴

SQLite作为嵌入式数据库，无需独立服务，完美契合桌面应用的数据存储需求。

4.2 数据库表设计

系统设计了两张表：`training_sessions`（会话表）和`epoch_metrics`（指标表）。

**会话表结构：**
 

CREATE TABLE training_sessions (
session_id INTEGER PRIMARY KEY AUTOINCREMENT,
start_time TEXT NOT NULL,
end_time TEXT,
data_yaml TEXT,
model_yaml TEXT,
pretrained_weights TEXT,
epochs INTEGER,
batch INTEGER,
imgsz INTEGER,
device TEXT,
optimizer TEXT,
notes TEXT
)

**指标表结构：**

CREATE TABLE epoch_metrics (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id INTEGER NOT NULL,
epoch INTEGER NOT NULL,
loss REAL,
box_loss REAL,
cls_loss REAL,
dfl_loss REAL,
metrics_json TEXT,
FOREIGN KEY (session_id) REFERENCES training_sessions (session_id) ON DELETE CASCADE
)

4.3 数据库操作封装

`MetricsDatabase`类封装了所有数据库操作：

class MetricsDatabase:
def __init__(self, db_path="training_metrics.db"):
self.db_path = db_path
self._init_db()

def create_session(self, start_time, config):
"""创建训练会话，返回session_id"""
# 插入会话记录
# 返回自增ID

def update_session_end(self, session_id, end_time):
"""更新会话结束时间"""

def save_epoch_metrics(self, session_id, epoch, metrics):
"""保存单个epoch的指标"""
# 将额外指标转为JSON存储

def get_all_sessions(self):
"""获取所有会话列表（用于历史记录下拉框）"""

def get_session_metrics(self, session_id):
"""获取指定会话的所有epoch指标"""

4.4 设计要点

• 外键约束：`epoch_metrics.session_id`关联到`training_sessions.session_id`，级联删除确保数据一致性。
• JSON字段：使用`metrics_json`存储非固定字段（如mAP、precision、recall等），保持表结构灵活。
• 时间戳格式：统一使用`YYYY-MM-DD HH:MM:SS`格式，便于排序和显示。
• 参数默认值：提供合理的默认值，避免NULL导致查询异常。

第5章 训练引擎：基于Ultralytics的异步训练

5.1 训练工作线程设计

GUI应用必须保持界面响应，因此训练任务必须在后台线程中执行。PyQt的`QThread`是实现这一模式的标准方式。

`TrainWorker`类设计如下：

class TrainWorker(QThread):
# 信号定义
log_signal = pyqtSignal(str) # 日志输出
progress_signal = pyqtSignal(int, int) # 进度更新
metric_signal = pyqtSignal(dict) # 指标更新
finished_signal = pyqtSignal(bool, str) # 完成信号
save_dir_signal = pyqtSignal(str) # 保存目录
session_created_signal = pyqtSignal(int) # 会话ID

def __init__(self, config):
super().__init__()
self.config = config
self.is_running = True

def run(self):
# 1. 初始化模型
# 2. 创建数据库会话
# 3. 注册回调
# 4. 执行训练
# 5. 更新会话结束时间

5.2 YOLO训练回调机制

Ultralytics YOLO支持在训练过程中注册回调函数。我们在`on_train_epoch_end`回调中提取指标：

def on_epoch_end(self, trainer):
# 提取总损失
tloss = getattr(trainer, 'tloss', None)
metrics['loss'] = to_scalar(tloss)

# 提取各项损失
if hasattr(trainer, 'loss_items'):
items = list(trainer.loss_items)
for i, key in enumerate(['box_loss', 'cls_loss', 'dfl_loss']):
if i < len(items):
metrics[key] = to_scalar(items[i])

# 提取其他指标（如mAP）
if hasattr(trainer, 'metrics'):
for k, v in trainer.metrics.items():
metrics[k] = to_scalar(v)

# 保存到数据库
self.db.save_epoch_metrics(self.session_id, epoch, metrics)

# 写入TensorBoard
if self.writer:
for k, v in metrics.items():
if k != 'epoch' and not isinstance(v, dict):
self.writer.add_scalar(k, float(v), epoch)

5.3 张量安全转换工具

训练过程中产生的指标可能是`torch.Tensor`、`numpy.ndarray`或Python原生类型。为了避免类型转换错误，我们实现了安全的转换函数：

def to_scalar(x):
"""安全地将各类数值转换为Python标量"""
if x is None:
return 0.0
if isinstance(x, torch.Tensor):
if x.numel() == 1:
return x.item()
else:
return float(x.mean().item())
if isinstance(x, np.ndarray):
if x.size == 1:
return x.item()
else:
return float(x.mean())
if isinstance(x, (int, float)):
return x
try:
return float(x)
except:
return 0.0

5.4 TensorBoard日志集成

TensorBoard是深度学习训练可视化的标准工具。我们在训练过程中手动写入日志，而不是依赖YOLO的自动日志（因为其存在字典类型处理bug）：

# 初始化TensorBoard writer
tb_dir = os.path.join(save_dir, "tensorboard")
self.writer = SummaryWriter(tb_dir)

# 每个epoch写入标量
for k, v in metrics.items():
if k == 'epoch' or isinstance(v, dict):
continue
try:
self.writer.add_scalar(k, float(v), epoch)
except (TypeError, ValueError):
pass

5.5 训练参数配置

训练参数分为三类，通过UI控件调节：

• 基础参数：数据集路径、模型配置、训练轮数、批次大小、图像尺寸、设备、优化器
• 超参数：学习率（初始/最终）、动量、权重衰减、预热参数、损失增益
• 数据增强：HSV变换、旋转、平移、缩放、剪切、透视、翻转、Mosaic、MixUp

所有参数最终组装为字典传入`model.train(**train_args)`。

第6章 实时指标与TensorBoard可视化

6.1 实时指标显示

训练过程中，主界面通过`metric_signal`接收指标更新，并在"实时指标"标签页显示：

def update_metrics_display(self, metrics):
epoch = metrics.get('epoch', 0)
loss = to_scalar(metrics.get('loss', 0.0))
box = to_scalar(metrics.get('box_loss', 0.0))
cls = to_scalar(metrics.get('cls_loss', 0.0))
dfl = to_scalar(metrics.get('dfl_loss', 0.0))
text = f"Epoch {epoch} | 总损失: {loss:.4f} | Box: {box:.4f} | Cls: {cls:.4f} | DFL: {dfl:.4f}"
self.metrics_label.setText(text)

这种设计让用户在不切换窗口的情况下快速掌握训练动态。

### 6.2 TensorBoard服务管理

系统启动时自动启动TensorBoard服务：

def init_tensorboard_server(self):
log_dir = "runs"
self.tb_process = subprocess.Popen(
["tensorboard", "–logdir", log_dir, "–port", "6006", "–bind_all"],
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
)
time.sleep(2)
self.tb_view.setUrl(QUrl("http://localhost:6006"))

使用`QWebEngineView`在GUI中嵌入TensorBoard界面，实现"一站式"体验。

### 6.3 定时刷新机制

为了确保TensorBoard页面及时更新，我们使用`QTimer`定时刷新：

self.tb_timer = QTimer()
self.tb_timer.timeout.connect(self.refresh_tensorboard)
self.tb_timer.start(5000) # 每5秒刷新

第7章 历史记录与曲线绘制

7.1 会话列表管理

`HistoryWidget`从数据库中加载所有训练会话，填充到下拉框中：

def load_session_list(self):
sessions = self.db.get_all_sessions()
for sess in sessions:
session_id = sess[0]
start_time = sess[1]
end_time = sess[2] if sess[2] else "进行中"
data_yaml = sess[3]
epochs = sess[4]
display = f"ID:{session_id} | {start_time} | {end_time} | {os.path.basename(data_yaml)} | {epochs} epochs"
self.session_combo.addItem(display, session_id)

7.2 Matplotlib中文字体解决方案

Matplotlib默认不支持中文，需要手动配置字体：

def set_chinese_font(self):
chinese_fonts = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Zen Hei', 'Noto Sans CJK SC']
available_fonts = [f.name for f in fm.fontManager.ttflist]
for font_name in chinese_fonts:
if font_name in available_fonts:
plt.rcParams['font.sans-serif'] = [font_name]
break
plt.rcParams['axes.unicode_minus'] = False

7.3 动态曲线绘制

用户可以选择显示哪些指标（总损失、Box损失、Cls损失、DFL损失），点击"绘制曲线"按钮生成图表：

def plot_metrics(self):
self.figure.clear()
ax = self.figure.add_subplot(111)
epochs = [m['epoch'] for m in self.metrics_data]

if self.loss_cb.isChecked():
losses = [m['loss'] for m in self.metrics_data]
ax.plot(epochs, losses, label='总损失', marker='o')

# 其他指标同理…
ax.set_xlabel("训练轮数 (Epoch)")
ax.set_ylabel("损失值 (Loss)")
ax.legend()
ax.grid(True)
self.canvas.draw()

7.4 图表保存功能

支持将当前曲线图保存为PNG、JPG或PDF格式，方便报告和文档使用。

第8章 模型测试与效果验证

8.1 智能路径解析

测试需要定位训练图片目录，系统支持三种方式：

• 自动解析：从`data.yaml`中读取`path`和`train`字段，组合出绝对路径
• 手动选择：用户通过文件对话框指定图片目录
• 智能回退：遍历`runs/train`和`runs/detect/train`下最新的`exp`目录

路径解析逻辑：

def get_train_image_dir(self):
data_yaml = self.data_edit.text()
with open(data_yaml, 'r') as f:
data_cfg = yaml.safe_load(f)

train_path = data_cfg.get('train')
base_path = data_cfg.get('path', '')

# 尝试多种组合方式
candidates = []
if train_path:
# 相对于data.yaml所在目录
candidates.append(os.path.join(os.path.dirname(data_yaml), train_path))
# 相对于path字段
if base_path:
candidates.append(os.path.join(base_path, train_path))
# 回退到常见目录结构
candidates.append(os.path.join(os.path.dirname(data_yaml), 'images', 'train'))

for cand in candidates:
if os.path.exists(cand):
return cand
return None

8.2 随机图片测试

从训练图片目录中随机抽取两张图片，分别用`best.pt`和`last.pt`进行推理：

def run_image_test(self, manual_best=None, manual_last=None):
# 1. 确定权重路径（自动/手动）
# 2. 获取图片目录
# 3. 随机抽取2张图片
# 4. 加载模型进行推理
# 5. 显示预测结果

8.3 结果可视化

将预测结果绘制在图片上，并垂直拼接展示：

def combine_images(self, results, img_paths, target_width=640):
imgs = []
for res, path in zip(results, img_paths):
img = res.plot() # YOLO内置绘制
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 添加文件名标注
cv2.putText(img_rgb, os.path.basename(path), (10, 30),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
imgs.append(img_rgb)

# 统一宽度并垂直拼接
resized = [cv2.resize(img, (target_width, int(img.shape[0] * target_width / img.shape[1])) for img in imgs]
concat = np.vstack(resized)
# 转为QPixmap显示

8.4 手动权重测试

支持用户通过"手动选择权重测试"按钮加载任意`.pt`文件进行推理，方便对比不同训练版本的效果。

第9章 模型剪枝：结构化通道剪枝

9.1 剪枝原理简介

结构化剪枝（Structured Pruning）通过删除整个卷积通道或神经元，直接减小模型宽度和计算量，相比非结构化剪枝（稀疏化）更有利于硬件加速。

torch_pruning库的实现思路：

1. 构建模型的计算图依赖关系
2. 根据重要性指标（如权重幅度）对通道排序
3. 按比例移除不重要的通道
4. 微调恢复精度（本系统暂未集成自动微调）

9.2 剪枝工作线程实现

class PruneWorker(QThread):
def run(self):
model = YOLO(self.weight_path).model
example_inputs = torch.randn(1, 3, 640, 640).to(device)

# 计算剪枝前FLOPs
self._log_flops(model, example_inputs, "剪枝前")

# 构建剪枝器
imp = tp.importance.MagnitudeImportance(p=2)
pruner = tp.pruner.MetaPruner(
model,
example_inputs,
importance=imp,
pruning_ratio=self.prune_ratio,
iterative_steps=1,
)

# 执行剪枝
pruner.step()

# 计算剪枝后FLOPs
self._log_flops(model, example_inputs, "剪枝后")

# 保存剪枝后权重
save_path = self.weight_path.replace('.pt', '_pruned.pt')
torch.save(model.state_dict(), save_path)

9.3 FLOPs估算

使用`thop`库估算模型计算量：

from thop import profile, clever_format
flops, _ = profile(model, inputs=(example_inputs,), verbose=False)
flops_g, _ = clever_format([flops, 0], "%.3f")
self.log_signal.emit(f"{stage_name} GFLOPs: {flops_g}")

9.4 剪枝实践建议

第10章 模型导出：ONNX与OpenVINO

10.1 导出格式选择

• ONNX：跨平台推理、移动端部署；通用性强，支持各种推理引擎 。
• OpenVINO ：Intel CPU/VPU/GPU加速；针对Intel硬件优化，支持INT8量化。
• TensorRT ：NVIDIA GPU推理；推理速度最快，但依赖NVIDIA硬件。
• CoreML ：Apple设备；iOS/macOS原生支持。

本系统实现了ONNX和OpenVINO两种常用格式的导出。

10.2 ONNX导出实现

def export_model(self, format_type):
model = YOLO(self.weight_path)
if format_type == "onnx":
model.export(
format="onnx",
imgsz=640,
simplify=True, # 简化计算图
opset=12, # ONNX算子集版本
nms=False # 不包含NMS层（减小体积）
)

10.3 OpenVINO导出实现

elif format_type == "openvino":
data_yaml = QFileDialog.getOpenFileName(…) # 可选，用于INT8量化
model.export(
format="openvino",
imgsz=640,
int8=bool(data_yaml), # 是否启用INT8量化
data=data_yaml if data_yaml else None
)

10.4 导出注意事项

动态尺寸：默认导出固定尺寸（640×640），如需动态尺寸需设置`dynamic=True`。
简化优化：`simplify=True`会执行常量折叠等优化，减小模型体积。
INT8量化：需要提供校准数据集（通过data.yaml指定），量化后模型体积减小约75%，推理速度提升2-4倍。
算子兼容性：注意目标推理框架支持的ONNX opset版本。

第11章 性能优化与最佳实践

11.1 内存管理优化

问题：训练过程中频繁创建张量和模型副本，可能导致内存泄漏。

解决方案：

1. 使用`torch.cuda.empty_cache()`定期释放显存
2. 在回调中避免持有大型张量的引用
3. 数据库批量插入减少事务开销

11.2 训练速度优化

• 增大batch size（显存允许），提高GPU利用率。
• 使用混合精度训练（amp），速度提升30-50%。
• 减少数据增强复杂度，降低CPU预处理开销。
• 使用DDP多卡训练，线性加速。

系统在`train_args`中默认未开启amp，可通过修改代码添加。

11.3 GUI响应优化

原则：耗时操作（训练、剪枝、导出）全部在子线程执行，主线程只负责UI更新和信号接收。

注意事项：
信号传递的数据应尽量小（避免传递大型numpy数组）
高频信号（如每epoch的指标）使用队列缓冲，避免UI卡顿

11.4 数据增强调优指南

|     增强参数       |        推荐值      |                                   说明                             |
|———————-|——————– |———————————————————– |
|         hsv_h        |         0.015       | 色调扰动，提高颜色鲁棒性                           |
|         hsv_s        |           0.7         | 饱和度扰动                                                    |
|        mosaic       |          1.0          | Mosaic增强，小目标检测效果好                   |
|         mixup        |        0.0~0.5     | 混合增强，防止过拟合                                  |
|    copy_paste    |        0.0~0.3     | 复制粘贴增强，适用于实例分割                    |

内存不足时可调低：mosaic → 0.5，imgsz → 512，batch → 4。

第12章 实战案例：训练一个自定义数据集

12.1 数据集准备

假设我们要训练一个"安全帽检测"模型，数据集结构如下：
helmet_dataset/
├── data.yaml
├── images/
│   ├── train/  # 5000张图片
│   └── val/    # 500张图片
└── labels/
    ├── train/  # 对应的YOLO格式标签
    └── val/

`data.yaml`:
path: ./helmet_dataset
train: images/train
val: images/val
nc: 2
names: ['helmet', 'person']
 

12.2 训练配置

在GUI中设置参数：

数据集YAML：`helmet_dataset/data.yaml`
模型YAML：`ultralytics/cfg/models/26/yolo26n.yaml`
预训练权重：`yolo26n.pt`
训练轮数：150
批次大小：16（显存8GB推荐）
图像尺寸：640
初始学习率：0.01

12.3 训练过程监控

训练过程中，可以观察到：

实时指标：标签页显示当前epoch的损失值
TensorBoard：显示完整的loss曲线和mAP曲线
进度条：展示训练进度

如下所示：

12.4 效果测试

训练完成后，系统自动执行随机图片测试，展示`best.pt`和`last.pt`的预测效果。

12.5 模型优化

在"模型管理"标签页：

1. 选择训练好的`best.pt`
2. 设置剪枝率为0.3
3. 执行剪枝，对比剪枝前后FLOPs
4. 导出ONNX格式用于部署

12.6 结果分析

通过"历史记录"标签页查看完整的训练曲线，对比不同会话的训练效果。

第13章 扩展与定制开发

13.1 支持更多YOLO版本

系统使用Ultralytics YOLO框架，理论上支持YOLOv5/v8/v9/v10/v11/v26等所有版本。只需修改`model_yaml`路径即可。

13.2 添加更多导出格式

在`ModelManagerWidget`中增加按钮，调用`model.export()`即可支持：

# TensorRT导出
model.export(format="engine", imgsz=640)
# CoreML导出
model.export(format="coreml", imgsz=640)

13.3 增加模型评估功能

可在测试标签页增加"评估"按钮，在验证集上计算mAP、Precision、Recall等指标。

13.4 增加自动超参数搜索

集成Optuna或Ray Tune，实现超参数自动优化。

13.5 云端训练支持

将训练任务提交到云端GPU集群，通过SSH或API进行远程训练。

第14章 常见问题与解决方案

14.1 训练启动失败

现象：点击"开始训练"后无反应或报错。

排查步骤：
1. 检查data.yaml是否存在且格式正确
2. 检查模型YAML路径是否正确
3. 检查预训练权重是否可读
4. 查看"训练日志"中的错误信息

14.2 显存不足

现象：训练过程中报CUDA out of memory。

解决方案：
1. 降低batch size（如从16降到8）
2. 降低imgsz（如从640降到512）
3. 关闭Mosaic增强（mosaic=0）
4. 减少workers数量

14.3 TensorBoard无法访问

现象：TensorBoard标签页空白或显示"无法访问此网站"。

解决方案：
1. 检查6006端口是否被占用：`netstat -ano | findstr 6006`
2. 杀死占用进程后重启系统
3. 手动启动TensorBoard：`tensorboard –logdir runs –port 6006`

14.4 剪枝失败

现象：剪枝过程中报错，提示模型结构不支持。

原因：torch_pruning对某些YOLO版本的支持有限。

解决方案：
1. 检查torch_pruning版本，建议使用1.1.0+
2. 尝试降低剪枝率
3. 更新至最新YOLO版本

14.5 数据库文件损坏

现象：历史记录无法加载或程序崩溃。

解决方案：
1. 备份`training_metrics.db`文件
2. 删除后重新运行程序（会自动创建新数据库）
3. 使用SQLite工具修复：`sqlite3 training_metrics.db "PRAGMA integrity_check;"`

第15章 总结

15.1 系统特点总结

本文构建的YOLO一体化平台具有以下特点：

15.2 技术亮点

• 多线程架构：训练、剪枝等耗时操作在后台执行，界面保持流畅。
• 信号驱动通信：使用PyQt信号在子线程和主线程间传递数据，线程安全。
• 数据库持久化：结构化存储训练指标，支持复杂的查询和对比分析。
• 回调机制：深度集成Ultralytics训练流程，精确控制每个epoch的数据采集。
• 智能路径解析：自动从data.yaml解析训练图片目录，减少用户配置。

15.3 结语

本文从架构设计到代码实现，完整介绍了YOLO一体化平台的搭建过程。该平台将训练、测试、优化、部署等环节整合到一个统一的GUI界面中，大幅提升了开发效率。

通过开源代码的分享，我们希望降低目标检测技术的使用门槛，让更多开发者能够专注于业务逻辑而非工具链的搭建。无论你是刚入门的深度学习爱好者，还是经验丰富的算法工程师，都能从这个平台中获得便利。

未来，随着YOLO版本的持续迭代和硬件加速技术的进步，这个平台也将不断演进，为更多场景提供智能化的目标检测解决方案。

附录A：完整代码清单

# -*- coding: utf-8 -*-
"""
YOLOv26 智能训练系统
– 支持超参数调节（学习率、数据增强等）
– 动态获取实际训练输出目录，自动定位 best.pt / last.pt
– 手动记录 TensorBoard 日志（解决字典类型错误）
– 智能解析 data.yaml 中的训练图片目录 + 手动选择备用
– 实时指标显示（解决张量转换错误）
– 随机图片测试 + 手动选择权重
– 训练指标数据库存储（SQLite），支持历史记录查询与 Matplotlib 曲线绘制
"""

import sys
import os
import csv
import time
import random
import subprocess
import sqlite3
import datetime
from pathlib import Path

import cv2
import numpy as np
import yaml
import torch
import torch_pruning as tp
from torch.utils.tensorboard import SummaryWriter

from PyQt5.QtWidgets import (
QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
QLabel, QPushButton, QLineEdit, QSpinBox, QDoubleSpinBox,
QComboBox, QTabWidget, QTextEdit, QProgressBar,
QMessageBox, QGroupBox, QGridLayout, QSplitter, QFileDialog,
QScrollArea, QFrame, QDateEdit, QDateTimeEdit, QCheckBox,
QListWidget, QListWidgetItem, QFormLayout
)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer, QUrl, QDateTime
from PyQt5.QtGui import QFont, QPixmap, QImage
from PyQt5.QtWebEngineWidgets import QWebEngineView

# Matplotlib 相关
import matplotlib
matplotlib.use('Qt5Agg')
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm

from ultralytics import YOLO

# ======================= 辅助函数：安全地将张量转为标量 =======================
def to_scalar(x):
"""将 torch.Tensor 或 numpy 数组安全地转换为 Python float / int"""
if x is None:
return 0.0
if isinstance(x, torch.Tensor):
if x.numel() == 1:
return x.item()
else:
return float(x.mean().item())
if isinstance(x, np.ndarray):
if x.size == 1:
return x.item()
else:
return float(x.mean())
if isinstance(x, (int, float)):
return x
try:
return float(x)
except:
return 0.0

# ========================== 数据库管理类 ==========================
class MetricsDatabase:
"""管理训练指标数据库（SQLite），存储会话和每个 epoch 的指标"""
def __init__(self, db_path="training_metrics.db"):
self.db_path = db_path
self._init_db()

def _init_db(self):
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
# 会话表：记录每次训练运行
cursor.execute('''
CREATE TABLE IF NOT EXISTS training_sessions (
session_id INTEGER PRIMARY KEY AUTOINCREMENT,
start_time TEXT NOT NULL,
end_time TEXT,
data_yaml TEXT,
model_yaml TEXT,
pretrained_weights TEXT,
epochs INTEGER,
batch INTEGER,
imgsz INTEGER,
device TEXT,
optimizer TEXT,
notes TEXT
)
''')
# 指标表：每个 epoch 的详细指标
cursor.execute('''
CREATE TABLE IF NOT EXISTS epoch_metrics (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id INTEGER NOT NULL,
epoch INTEGER NOT NULL,
loss REAL,
box_loss REAL,
cls_loss REAL,
dfl_loss REAL,
metrics_json TEXT,
FOREIGN KEY (session_id) REFERENCES training_sessions (session_id) ON DELETE CASCADE
)
''')
conn.commit()
conn.close()

def create_session(self, start_time, config):
"""创建新的训练会话，返回 session_id"""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute('''
INSERT INTO training_sessions
(start_time, data_yaml, model_yaml, pretrained_weights, epochs, batch, imgsz, device, optimizer, notes)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
''', (
start_time,
config.get('data', ''),
config.get('model_yaml', ''),
config.get('pretrained_weights', ''),
config.get('epochs', 0),
config.get('batch', 0),
config.get('imgsz', 0),
config.get('device', ''),
config.get('optimizer', ''),
config.get('notes', '')
))
session_id = cursor.lastrowid
conn.commit()
conn.close()
return session_id

def update_session_end(self, session_id, end_time):
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute('UPDATE training_sessions SET end_time = ? WHERE session_id = ?', (end_time, session_id))
conn.commit()
conn.close()

def save_epoch_metrics(self, session_id, epoch, metrics):
"""保存单个 epoch 的指标"""
import json
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
# 将额外指标转为 JSON 字符串
metrics_json = json.dumps({k: v for k, v in metrics.items() if k not in ['epoch', 'loss', 'box_loss', 'cls_loss', 'dfl_loss']})
cursor.execute('''
INSERT INTO epoch_metrics (session_id, epoch, loss, box_loss, cls_loss, dfl_loss, metrics_json)
VALUES (?, ?, ?, ?, ?, ?, ?)
''', (
session_id, epoch,
metrics.get('loss', 0.0),
metrics.get('box_loss', 0.0),
metrics.get('cls_loss', 0.0),
metrics.get('dfl_loss', 0.0),
metrics_json
))
conn.commit()
conn.close()

def get_all_sessions(self):
"""获取所有会话，返回列表 [(session_id, start_time, end_time, …)]"""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute('SELECT session_id, start_time, end_time, data_yaml, epochs FROM training_sessions ORDER BY start_time DESC')
rows = cursor.fetchall()
conn.close()
return rows

def get_session_metrics(self, session_id):
"""获取指定会话的所有 epoch 指标"""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute('SELECT epoch, loss, box_loss, cls_loss, dfl_loss, metrics_json FROM epoch_metrics WHERE session_id = ? ORDER BY epoch', (session_id,))
rows = cursor.fetchall()
conn.close()
# 转换为字典列表
metrics_list = []
for row in rows:
metrics_list.append({
'epoch': row[0],
'loss': row[1],
'box_loss': row[2],
'cls_loss': row[3],
'dfl_loss': row[4],
'metrics_json': row[5]
})
return metrics_list

# ========================== 训练工作线程 ==========================
class TrainWorker(QThread):
log_signal = pyqtSignal(str)
progress_signal = pyqtSignal(int, int)
metric_signal = pyqtSignal(dict)
finished_signal = pyqtSignal(bool, str)
save_dir_signal = pyqtSignal(str) # 发送实际保存目录
session_created_signal = pyqtSignal(int) # 发送 session_id

def __init__(self, config):
super().__init__()
self.config = config
self.model = None
self.is_running = True
self.epoch = 0
self.total_epochs = config['epochs']
self.save_dir_sent = False
self.writer = None
self.db = MetricsDatabase()
self.session_id = None

def run(self):
try:
self.log_signal.emit("[INFO] 初始化模型…")
self.model = YOLO(self.config['model_yaml'])
if self.config.get('pretrained_weights') and os.path.exists(self.config['pretrained_weights']):
self.model.load(self.config['pretrained_weights'])
self.log_signal.emit(f"[INFO] 加载预训练权重: {self.config['pretrained_weights']}")

# 创建数据库会话记录
start_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
self.session_id = self.db.create_session(start_time, self.config)
self.session_created_signal.emit(self.session_id)
self.log_signal.emit(f"[INFO] 训练会话已记录，ID: {self.session_id}")

# 注册自定义回调
self.model.add_callback("on_train_epoch_end", self.on_epoch_end)

# 准备训练参数字典
train_args = {
'data': self.config['data'],
'epochs': self.config['epochs'],
'batch': self.config['batch'],
'imgsz': self.config['imgsz'],
'device': self.config['device'],
'workers': self.config['workers'],
'project': self.config['project'],
'name': self.config['name'],
'exist_ok': True,
'pretrained': self.config.get('pretrained', True),
'optimizer': self.config['optimizer'],
'lr0': self.config['lr0'],
'lrf': self.config['lrf'],
'momentum': self.config['momentum'],
'weight_decay': self.config['weight_decay'],
'warmup_epochs': self.config['warmup_epochs'],
'warmup_momentum': self.config['warmup_momentum'],
'warmup_bias_lr': self.config['warmup_bias_lr'],
'box': self.config['box_loss_gain'],
'cls': self.config['cls_loss_gain'],
'dfl': self.config['dfl_loss_gain'],
'hsv_h': self.config['hsv_h'],
'hsv_s': self.config['hsv_s'],
'hsv_v': self.config['hsv_v'],
'degrees': self.config['degrees'],
'translate': self.config['translate'],
'scale': self.config['scale'],
'shear': self.config['shear'],
'perspective': self.config['perspective'],
'flipud': self.config['flipud'],
'fliplr': self.config['fliplr'],
'mosaic': self.config['mosaic'],
'mixup': self.config['mixup'],
'copy_paste': self.config['copy_paste'],
'close_mosaic': self.config['close_mosaic'],
}
train_args = {k: v for k, v in train_args.items() if v is not None}

self.log_signal.emit("[INFO] 开始训练…")
results = self.model.train(**train_args)
if self.writer:
self.writer.close()
# 更新会话结束时间
end_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
self.db.update_session_end(self.session_id, end_time)
self.finished_signal.emit(True, "训练完成！")
except Exception as e:
if self.writer:
self.writer.close()
self.finished_signal.emit(False, f"训练出错: {str(e)}")

def on_epoch_end(self, trainer):
if not self.is_running:
trainer.stop = True
return
self.epoch = trainer.epoch
self.progress_signal.emit(self.epoch, self.total_epochs)

# 第一次回调时发送保存目录并初始化 TensorBoard writer
if not self.save_dir_sent:
save_dir = str(trainer.save_dir)
self.save_dir_signal.emit(save_dir)
self.save_dir_sent = True
tb_dir = os.path.join(save_dir, "tensorboard")
os.makedirs(tb_dir, exist_ok=True)
self.writer = SummaryWriter(tb_dir)
self.log_signal.emit(f"[INFO] TensorBoard 日志将保存至: {tb_dir}")

# 提取指标
metrics = {'epoch': self.epoch}
tloss = getattr(trainer, 'tloss', None)
metrics['loss'] = to_scalar(tloss)

if hasattr(trainer, 'loss_items') and trainer.loss_items is not None:
items = list(trainer.loss_items) if hasattr(trainer.loss_items, '__len__') else [trainer.loss_items]
for i, key in enumerate(['box_loss', 'cls_loss', 'dfl_loss']):
if i < len(items):
metrics[key] = to_scalar(items[i])

if hasattr(trainer, 'metrics'):
for k, v in trainer.metrics.items():
metrics[k] = to_scalar(v)

# 保存到数据库
if self.session_id is not None:
self.db.save_epoch_metrics(self.session_id, self.epoch, metrics)

# 记录到 TensorBoard（跳过非数值类型，如字典）
if self.writer:
step = self.epoch
for k, v in metrics.items():
if k == 'epoch':
continue
if isinstance(v, dict):
continue
try:
scalar_val = float(v)
self.writer.add_scalar(k, scalar_val, step)
except (TypeError, ValueError):
pass
if hasattr(trainer, 'lr'):
lr = trainer.lr
try:
if isinstance(lr, (int, float)):
self.writer.add_scalar('lr', lr, step)
elif isinstance(lr, torch.Tensor) and lr.numel() == 1:
self.writer.add_scalar('lr', lr.item(), step)
except:
pass

self.metric_signal.emit(metrics)

# 写入 CSV（保留原有逻辑）
csv_path = os.path.join(trainer.save_dir, 'training_metrics.csv')
write_header = not os.path.exists(csv_path)
with open(csv_path, 'a', newline='') as f:
writer = csv.DictWriter(f, fieldnames=metrics.keys())
if write_header:
writer.writeheader()
writer.writerow(metrics)

def stop(self):
self.is_running = False
if self.model:
self.model.stop = True

# ========================== 历史记录标签页：Matplotlib 曲线 ==========================
class HistoryWidget(QWidget):
def __init__(self, parent=None):
super().__init__(parent)
self.db = MetricsDatabase()
self.current_session_id = None
self.metrics_data = [] # 存储当前选中会话的指标列表
self.init_ui()
self.set_chinese_font() # 设置中文字体

def set_chinese_font(self):
"""自动检测并设置 matplotlib 中文字体"""
# 列出常见中文字体名称（按优先级排序）
chinese_fonts = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Zen Hei', 'Noto Sans CJK SC', 'STHeiti', 'PingFang SC']
available_fonts = [f.name for f in fm.fontManager.ttflist]
for font_name in chinese_fonts:
if font_name in available_fonts:
plt.rcParams['font.sans-serif'] = [font_name]
break
else:
# 如果没有找到中文字体，尝试使用系统默认中文字体（可能不是完美匹配）
# 或提示用户
print("警告: 未找到中文字体，图表中的中文将显示为方块。请安装 SimHei 或 Microsoft YaHei 字体。")
# 仍尝试使用默认字体，避免报错
plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题

def init_ui(self):
layout = QVBoxLayout(self)

# 上方控制区
control_layout = QHBoxLayout()
self.session_combo = QComboBox()
self.session_combo.setMinimumWidth(400)
self.session_combo.currentIndexChanged.connect(self.on_session_changed)
control_layout.addWidget(QLabel("选择训练会话:"))
control_layout.addWidget(self.session_combo)

self.refresh_btn = QPushButton("刷新列表")
self.refresh_btn.clicked.connect(self.load_session_list)
control_layout.addWidget(self.refresh_btn)

layout.addLayout(control_layout)

# 指标选择区域（多选框）
metrics_group = QGroupBox("选择要显示的指标")
metrics_layout = QHBoxLayout(metrics_group)
self.loss_cb = QCheckBox("总损失 (loss)")
self.loss_cb.setChecked(True)
self.box_loss_cb = QCheckBox("Box 损失")
self.box_loss_cb.setChecked(True)
self.cls_loss_cb = QCheckBox("Cls 损失")
self.cls_loss_cb.setChecked(True)
self.dfl_loss_cb = QCheckBox("DFL 损失")
self.dfl_loss_cb.setChecked(True)
# 可以根据需要添加更多
metrics_layout.addWidget(self.loss_cb)
metrics_layout.addWidget(self.box_loss_cb)
metrics_layout.addWidget(self.cls_loss_cb)
metrics_layout.addWidget(self.dfl_loss_cb)
metrics_layout.addStretch()
layout.addWidget(metrics_group)

# Matplotlib 图形区域
self.figure = Figure(figsize=(10, 6), dpi=100)
self.canvas = FigureCanvas(self.figure)
layout.addWidget(self.canvas)

# 底部按钮区
btn_layout = QHBoxLayout()
self.plot_btn = QPushButton("绘制曲线")
self.plot_btn.clicked.connect(self.plot_metrics)
self.save_btn = QPushButton("保存图片")
self.save_btn.clicked.connect(self.save_figure)
btn_layout.addWidget(self.plot_btn)
btn_layout.addWidget(self.save_btn)
btn_layout.addStretch()
layout.addLayout(btn_layout)

# 初始加载会话列表
self.load_session_list()

def load_session_list(self):
"""从数据库加载所有会话，更新下拉框"""
sessions = self.db.get_all_sessions()
self.session_combo.clear()
self.session_data = [] # 存储每个条目的 (session_id, display_text)
for sess in sessions:
session_id = sess[0]
start_time = sess[1]
end_time = sess[2] if sess[2] else "进行中"
data_yaml = sess[3]
epochs = sess[4]
display = f"ID:{session_id} | {start_time} | {end_time} | {os.path.basename(data_yaml)} | {epochs} epochs"
self.session_combo.addItem(display, session_id)
self.session_data.append((session_id, display))
if self.session_combo.count() > 0:
self.on_session_changed(0)

def on_session_changed(self, index):
if index < 0:
return
session_id = self.session_combo.itemData(index)
self.current_session_id = session_id
# 加载该会话的指标数据
self.metrics_data = self.db.get_session_metrics(session_id)
# 自动绘制曲线
self.plot_metrics()

def plot_metrics(self):
"""根据选择的指标绘制曲线"""
if not self.metrics_data:
self.figure.clear()
ax = self.figure.add_subplot(111)
ax.text(0.5, 0.5, "无数据", transform=ax.transAxes, ha='center', va='center')
self.canvas.draw()
return

self.figure.clear()
ax = self.figure.add_subplot(111)

epochs = [m['epoch'] for m in self.metrics_data]
# 注意：epoch 可能不连续，但按顺序绘制即可
if self.loss_cb.isChecked():
losses = [m['loss'] for m in self.metrics_data]
ax.plot(epochs, losses, label='总损失', marker='o', linestyle='-')
if self.box_loss_cb.isChecked():
box_losses = [m['box_loss'] for m in self.metrics_data]
ax.plot(epochs, box_losses, label='Box 损失', marker='s', linestyle='–')
if self.cls_loss_cb.isChecked():
cls_losses = [m['cls_loss'] for m in self.metrics_data]
ax.plot(epochs, cls_losses, label='Cls 损失', marker='^', linestyle='-.')
if self.dfl_loss_cb.isChecked():
dfl_losses = [m['dfl_loss'] for m in self.metrics_data]
ax.plot(epochs, dfl_losses, label='DFL 损失', marker='d', linestyle=':')

ax.set_xlabel("训练轮数 (Epoch)")
ax.set_ylabel("损失值 (Loss)")
ax.set_title(f"训练指标曲线 (会话 ID: {self.current_session_id})")
ax.legend()
ax.grid(True)
self.figure.tight_layout()
self.canvas.draw()

def save_figure(self):
"""保存当前曲线图为图片文件"""
if not self.metrics_data:
QMessageBox.warning(self, "提示", "没有数据可保存")
return
file_path, _ = QFileDialog.getSaveFileName(self, "保存图片", "", "PNG Image (*.png);;JPEG Image (*.jpg);;PDF File (*.pdf)")
if file_path:
self.figure.savefig(file_path, dpi=300, bbox_inches='tight')
QMessageBox.information(self, "成功", f"图片已保存至: {file_path}")
# ========================== 新增标签页 ==========================
class ModelManagerWidget(QWidget):
def __init__(self):
super().__init__()
self.weight_path = None
self.init_ui()

def init_ui(self):
layout = QVBoxLayout(self)

# 权重选择
weight_group = QGroupBox("模型选择")
weight_layout = QHBoxLayout(weight_group)
self.weight_line = QLineEdit()
self.weight_line.setPlaceholderText("请选择需要剪枝或导出的 .pt 权重文件")
weight_layout.addWidget(self.weight_line)
self.browse_btn = QPushButton("选择权重")
self.browse_btn.clicked.connect(self.select_weight_file)
weight_layout.addWidget(self.browse_btn)
layout.addWidget(weight_group)

# 剪枝
prune_group = QGroupBox("模型剪枝 (torch_pruning)")
prune_layout = QHBoxLayout(prune_group)
prune_layout.addWidget(QLabel("剪枝率 (通道剔除比例):"))
self.prune_ratio_spin = QDoubleSpinBox()
self.prune_ratio_spin.setRange(0.1, 0.9)
self.prune_ratio_spin.setSingleStep(0.05)
self.prune_ratio_spin.setValue(0.3)
prune_layout.addWidget(self.prune_ratio_spin)
self.start_prune_btn = QPushButton("开始结构化剪枝")
self.start_prune_btn.clicked.connect(self.start_pruning)
prune_layout.addWidget(self.start_prune_btn)
layout.addWidget(prune_group)

# 导出
export_group = QGroupBox("模型导出")
export_layout = QHBoxLayout(export_group)
self.export_onnx_btn = QPushButton("导出 ONNX")
self.export_onnx_btn.clicked.connect(lambda: self.export_model("onnx"))
self.export_openvino_btn = QPushButton("导出 OpenVINO IR")
self.export_openvino_btn.clicked.connect(lambda: self.export_model("openvino"))
export_layout.addWidget(self.export_onnx_btn)
export_layout.addWidget(self.export_openvino_btn)
layout.addWidget(export_group)

# 日志
log_group = QGroupBox("执行日志")
log_layout = QVBoxLayout(log_group)
self.log_text = QTextEdit()
self.log_text.setReadOnly(True)
log_layout.addWidget(self.log_text)
layout.addWidget(log_group)

def select_weight_file(self):
file_path, _ = QFileDialog.getOpenFileName(self, "选择权重文件", "", "PyTorch Model (*.pt)")
if file_path:
self.weight_path = file_path
self.weight_line.setText(file_path)

def start_pruning(self):
if not self.weight_path or not os.path.exists(self.weight_path):
QMessageBox.warning(self, "提示", "请先选择有效的权重文件")
return
self.start_prune_btn.setEnabled(False)
self.log_text.clear()
self.append_log(f"开始剪枝: {self.weight_path}, 剪枝率: {self.prune_ratio_spin.value()}")
self.prune_worker = PruneWorker(self.weight_path, self.prune_ratio_spin.value())
self.prune_worker.log_signal.connect(self.append_log)
self.prune_worker.finished_signal.connect(self.prune_finished)
self.prune_worker.start()

def prune_finished(self, success, msg):
self.start_prune_btn.setEnabled(True)
if success:
self.append_log(f"✅ 剪枝成功: {msg}")
QMessageBox.information(self, "完成", msg)
else:
self.append_log(f"❌ 剪枝失败: {msg}")
QMessageBox.critical(self, "错误", msg)

def append_log(self, msg):
self.log_text.append(msg)
self.log_text.ensureCursorVisible()

def export_model(self, format_type):
if not self.weight_path or not os.path.exists(self.weight_path):
QMessageBox.warning(self, "提示", "请先选择有效的权重文件")
return

save_dir = QFileDialog.getExistingDirectory(self, f"选择导出 {format_type.upper()} 模型保存目录")
if not save_dir:
return

self.append_log(f"开始导出 {format_type.upper()} 模型…")

try:
model = YOLO(self.weight_path)
if format_type == "onnx":
output_path = os.path.join(save_dir, "exported_model.onnx")
# 导出 ONNX，关闭 NMS 以减小体积
model.export(format="onnx", imgsz=640, simplify=True, opset=12, nms=False)
# 移动文件到目标目录
default_onnx = self.weight_path.replace('.pt', '.onnx')
if os.path.exists(default_onnx):
import shutil
shutil.move(default_onnx, output_path)
else:
# 尝试查找最新导出的 ONNX
import glob
onnx_files = glob.glob(os.path.join(os.path.dirname(self.weight_path), "*.onnx"))
if onnx_files:
shutil.move(onnx_files[0], output_path)
self.append_log(f"✅ ONNX 模型导出成功: {output_path}")

elif format_type == "openvino":
# 手动选择数据集 YAML（用于 INT8 量化，可选）
data_yaml, _ = QFileDialog.getOpenFileName(self, "选择数据集 YAML 文件 (用于 INT8 量化，可选)", "", "YAML files (*.yaml)")
if not data_yaml:
self.append_log("未选择数据集，将导出 FP32 格式 OpenVINO 模型")
# 导出 OpenVINO，如果提供了 data_yaml 则启用 int8 量化
model.export(format="openvino", imgsz=640, int8=bool(data_yaml), data=data_yaml if data_yaml else None)
self.append_log(f"✅ OpenVINO IR 模型导出成功，保存在: {os.path.dirname(self.weight_path)}")
QMessageBox.information(self, "完成", f"OpenVINO IR 模型已导出至 {os.path.dirname(self.weight_path)}")

except Exception as e:
self.append_log(f"❌ 导出失败: {str(e)}")
QMessageBox.critical(self, "错误", f"导出失败: {str(e)}")
# ========================== 剪枝工作线程 ==========================
class PruneWorker(QThread):
log_signal = pyqtSignal(str)
finished_signal = pyqtSignal(bool, str)

def __init__(self, weight_path, prune_ratio):
super().__init__()
self.weight_path = weight_path
self.prune_ratio = prune_ratio

def run(self):
try:
self.log_signal.emit("正在加载模型…")
model = YOLO(self.weight_path).model
example_inputs = torch.randn(1, 3, 640, 640).to(next(model.parameters()).device)

self.log_signal.emit("计算剪枝前模型计算量…")
self._log_flops(model, example_inputs, "剪枝前")

self.log_signal.emit("构建依赖图并执行剪枝…")
imp = tp.importance.MagnitudeImportance(p=2)
# 使用新版 torch_pruning 推荐参数
pruner = tp.pruner.MetaPruner(
model,
example_inputs,
importance=imp,
pruning_ratio=self.prune_ratio,
iterative_steps=1,
)
pruner.step()
self.log_signal.emit("剪枝完成。")

self.log_signal.emit("计算剪枝后模型计算量…")
self._log_flops(model, example_inputs, "剪枝后")

save_path = self.weight_path.replace('.pt', '_pruned.pt')
torch.save(model.state_dict(), save_path)
self.finished_signal.emit(True, f"剪枝后模型已保存至: {save_path}")
except Exception as e:
self.finished_signal.emit(False, str(e))

def _log_flops(self, model, example_inputs, stage_name):
try:
from thop import profile, clever_format
flops, _ = profile(model, inputs=(example_inputs,), verbose=False)
flops_g, _ = clever_format([flops, 0], "%.3f")
self.log_signal.emit(f"{stage_name} GFLOPs: {flops_g}")
except ImportError:
self.log_signal.emit(f"{stage_name} GFLOPs: 需要安装 thop (pip install thop) 以进行估算。")
except Exception as e:
self.log_signal.emit(f"{stage_name} GFLOPs 估算失败: {str(e)}")
# ========================== 主窗口 ==========================
class YOLOTrainGUI(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("YOLOv26 智能训练系统")
self.setGeometry(50, 50, 1500, 950)
self.setMinimumSize(1200, 700)
self.setStyleSheet(self.get_style())

self.train_worker = None
self.tb_process = None
self.tb_url = "http://localhost:6006"
self.train_dir = None
self.best_pt_path = None
self.last_pt_path = None
self.manual_image_dir = None

self.init_ui()
self.init_tensorboard_server()

# ———————— 样式表 ————————
def get_style(self):
return """
QMainWindow { background-color: #f0f2f5; }
QGroupBox {
font-weight: bold;
border: 1px solid #c0c4cc;
border-radius: 6px;
margin-top: 12px;
padding-top: 8px;
background-color: white;
}
QGroupBox::title {
subcontrol-origin: margin;
left: 12px;
padding: 0 5px 0 5px;
color: #2c3e50;
}
QLineEdit, QSpinBox, QDoubleSpinBox, QComboBox, QTextEdit {
border: 1px solid #dcdfe6;
border-radius: 4px;
padding: 4px;
background-color: white;
}
QPushButton {
background-color: #409eff;
color: white;
border: none;
border-radius: 4px;
padding: 6px 12px;
font-weight: bold;
}
QPushButton:hover { background-color: #66b1ff; }
QPushButton:pressed { background-color: #3a8ee6; }
QPushButton:disabled { background-color: #a6c9ff; }
QProgressBar {
border: 1px solid #dcdfe6;
border-radius: 4px;
text-align: center;
}
QProgressBar::chunk { background-color: #67c23a; }
QTabWidget::pane {
border: 1px solid #dcdfe6;
border-radius: 4px;
background-color: white;
}
QTabBar::tab {
background-color: #f5f7fa;
padding: 6px 12px;
margin-right: 2px;
}
QTabBar::tab:selected {
background-color: #409eff;
color: white;
}
"""

# ———————— UI 布局 ————————
def init_ui(self):
central = QWidget()
self.setCentralWidget(central)
main_layout = QVBoxLayout(central)
main_layout.setSpacing(10)
main_layout.setContentsMargins(10, 10, 10, 10)

# 参数区域（滚动）
param_scroll = QScrollArea()
param_scroll.setWidgetResizable(True)
param_scroll.setFrameShape(QFrame.NoFrame)
param_widget = QWidget()
param_layout = QVBoxLayout(param_widget)
param_layout.setSpacing(10)

# 基础参数
basic_group = QGroupBox("基础参数")
basic_grid = QGridLayout(basic_group)
self.data_edit = QLineEdit("data/data.yaml")
self.model_yaml_edit = QLineEdit("ultralytics/cfg/models/26/yolo26n.yaml")
self.weights_edit = QLineEdit("yolo26n.pt")
self.epochs_spin = QSpinBox(); self.epochs_spin.setRange(1,1000); self.epochs_spin.setValue(200)
self.batch_spin = QSpinBox(); self.batch_spin.setRange(1,128); self.batch_spin.setValue(8)
self.imgsz_spin = QSpinBox(); self.imgsz_spin.setRange(32,1280); self.imgsz_spin.setValue(640)
self.device_edit = QLineEdit("0")
self.workers_spin = QSpinBox(); self.workers_spin.setRange(0,16); self.workers_spin.setValue(2)
self.optimizer_combo = QComboBox(); self.optimizer_combo.addItems(["auto", "SGD", "Adam", "AdamW"])

row = 0
basic_grid.addWidget(QLabel("数据集YAML:"), row, 0); basic_grid.addWidget(self.data_edit, row, 1); row+=1
basic_grid.addWidget(QLabel("模型YAML:"), row, 0); basic_grid.addWidget(self.model_yaml_edit, row, 1); row+=1
basic_grid.addWidget(QLabel("预训练权重:"), row, 0); basic_grid.addWidget(self.weights_edit, row, 1); row+=1
basic_grid.addWidget(QLabel("训练轮数:"), row, 0); basic_grid.addWidget(self.epochs_spin, row, 1); row+=1
basic_grid.addWidget(QLabel("批次大小:"), row, 0); basic_grid.addWidget(self.batch_spin, row, 1); row+=1
basic_grid.addWidget(QLabel("图像尺寸:"), row, 0); basic_grid.addWidget(self.imgsz_spin, row, 1); row+=1
basic_grid.addWidget(QLabel("设备 (0/cpu):"), row, 0); basic_grid.addWidget(self.device_edit, row, 1); row+=1
basic_grid.addWidget(QLabel("工作线程数:"), row, 0); basic_grid.addWidget(self.workers_spin, row, 1); row+=1
basic_grid.addWidget(QLabel("优化器:"), row, 0); basic_grid.addWidget(self.optimizer_combo, row, 1)

param_layout.addWidget(basic_group)

# 超参数
hyper_group = QGroupBox("超参数（可调）")
hyper_layout = QGridLayout(hyper_group)
self.lr0_spin = QDoubleSpinBox(); self.lr0_spin.setRange(0,1); self.lr0_spin.setValue(0.01)
self.lrf_spin = QDoubleSpinBox(); self.lrf_spin.setRange(0,1); self.lrf_spin.setValue(0.01)
self.momentum_spin = QDoubleSpinBox(); self.momentum_spin.setRange(0,1); self.momentum_spin.setValue(0.937)
self.weight_decay_spin = QDoubleSpinBox(); self.weight_decay_spin.setRange(0,0.001); self.weight_decay_spin.setDecimals(5); self.weight_decay_spin.setValue(0.0005)
self.warmup_epochs_spin = QDoubleSpinBox(); self.warmup_epochs_spin.setRange(0,10); self.warmup_epochs_spin.setValue(3.0)
self.warmup_momentum_spin = QDoubleSpinBox(); self.warmup_momentum_spin.setRange(0,1); self.warmup_momentum_spin.setValue(0.8)
self.warmup_bias_lr_spin = QDoubleSpinBox(); self.warmup_bias_lr_spin.setRange(0,0.2); self.warmup_bias_lr_spin.setValue(0.1)
self.box_loss_spin = QDoubleSpinBox(); self.box_loss_spin.setRange(0,10); self.box_loss_spin.setValue(7.5)
self.cls_loss_spin = QDoubleSpinBox(); self.cls_loss_spin.setRange(0,10); self.cls_loss_spin.setValue(0.5)
self.dfl_loss_spin = QDoubleSpinBox(); self.dfl_loss_spin.setRange(0,10); self.dfl_loss_spin.setValue(1.5)

hyper_layout.addWidget(QLabel("初始学习率 lr0:"), 0,0); hyper_layout.addWidget(self.lr0_spin, 0,1)
hyper_layout.addWidget(QLabel("最终学习率 lrf:"), 1,0); hyper_layout.addWidget(self.lrf_spin, 1,1)
hyper_layout.addWidget(QLabel("动量 momentum:"), 2,0); hyper_layout.addWidget(self.momentum_spin, 2,1)
hyper_layout.addWidget(QLabel("权重衰减 weight_decay:"), 3,0); hyper_layout.addWidget(self.weight_decay_spin, 3,1)
hyper_layout.addWidget(QLabel("预热轮次 warmup_epochs:"), 4,0); hyper_layout.addWidget(self.warmup_epochs_spin, 4,1)
hyper_layout.addWidget(QLabel("预热动量 warmup_momentum:"), 5,0); hyper_layout.addWidget(self.warmup_momentum_spin, 5,1)
hyper_layout.addWidget(QLabel("预热偏置学习率:"), 6,0); hyper_layout.addWidget(self.warmup_bias_lr_spin, 6,1)
hyper_layout.addWidget(QLabel("box 损失增益:"), 7,0); hyper_layout.addWidget(self.box_loss_spin, 7,1)
hyper_layout.addWidget(QLabel("cls 损失增益:"), 8,0); hyper_layout.addWidget(self.cls_loss_spin, 8,1)
hyper_layout.addWidget(QLabel("dfl 损失增益:"), 9,0); hyper_layout.addWidget(self.dfl_loss_spin, 9,1)

param_layout.addWidget(hyper_group)

# 数据增强
aug_group = QGroupBox("数据增强（内存不足时可调低 mosaic/mixup 或降低 imgsz/batch）")
aug_layout = QGridLayout(aug_group)
self.hsv_h_spin = QDoubleSpinBox(); self.hsv_h_spin.setRange(0,1); self.hsv_h_spin.setValue(0.015)
self.hsv_s_spin = QDoubleSpinBox(); self.hsv_s_spin.setRange(0,1); self.hsv_s_spin.setValue(0.7)
self.hsv_v_spin = QDoubleSpinBox(); self.hsv_v_spin.setRange(0,1); self.hsv_v_spin.setValue(0.4)
self.degrees_spin = QDoubleSpinBox(); self.degrees_spin.setRange(0,180); self.degrees_spin.setValue(0.0)
self.translate_spin = QDoubleSpinBox(); self.translate_spin.setRange(0,1); self.translate_spin.setValue(0.1)
self.scale_spin = QDoubleSpinBox(); self.scale_spin.setRange(0,10); self.scale_spin.setValue(0.5)
self.shear_spin = QDoubleSpinBox(); self.shear_spin.setRange(0,90); self.shear_spin.setValue(0.0)
self.perspective_spin = QDoubleSpinBox(); self.perspective_spin.setRange(0,0.1); self.perspective_spin.setSingleStep(0.001); self.perspective_spin.setValue(0.0)
self.flipud_spin = QDoubleSpinBox(); self.flipud_spin.setRange(0,1); self.flipud_spin.setValue(0.0)
self.fliplr_spin = QDoubleSpinBox(); self.fliplr_spin.setRange(0,1); self.fliplr_spin.setValue(0.5)
self.mosaic_spin = QDoubleSpinBox(); self.mosaic_spin.setRange(0,1); self.mosaic_spin.setValue(1.0)
self.mixup_spin = QDoubleSpinBox(); self.mixup_spin.setRange(0,1); self.mixup_spin.setValue(0.0)
self.copy_paste_spin = QDoubleSpinBox(); self.copy_paste_spin.setRange(0,1); self.copy_paste_spin.setValue(0.0)
self.close_mosaic_spin = QSpinBox(); self.close_mosaic_spin.setRange(0,100); self.close_mosaic_spin.setValue(10)

aug_layout.addWidget(QLabel("HSV-H:"),0,0); aug_layout.addWidget(self.hsv_h_spin,0,1)
aug_layout.addWidget(QLabel("HSV-S:"),1,0); aug_layout.addWidget(self.hsv_s_spin,1,1)
aug_layout.addWidget(QLabel("HSV-V:"),2,0); aug_layout.addWidget(self.hsv_v_spin,2,1)
aug_layout.addWidget(QLabel("旋转角度:"),3,0); aug_layout.addWidget(self.degrees_spin,3,1)
aug_layout.addWidget(QLabel("平移:"),4,0); aug_layout.addWidget(self.translate_spin,4,1)
aug_layout.addWidget(QLabel("缩放:"),5,0); aug_layout.addWidget(self.scale_spin,5,1)
aug_layout.addWidget(QLabel("剪切:"),6,0); aug_layout.addWidget(self.shear_spin,6,1)
aug_layout.addWidget(QLabel("透视:"),7,0); aug_layout.addWidget(self.perspective_spin,7,1)
aug_layout.addWidget(QLabel("上下翻转概率:"),8,0); aug_layout.addWidget(self.flipud_spin,8,1)
aug_layout.addWidget(QLabel("左右翻转概率:"),9,0); aug_layout.addWidget(self.fliplr_spin,9,1)
aug_layout.addWidget(QLabel("Mosaic 概率:"),10,0); aug_layout.addWidget(self.mosaic_spin,10,1)
aug_layout.addWidget(QLabel("MixUp 概率:"),11,0); aug_layout.addWidget(self.mixup_spin,11,1)
aug_layout.addWidget(QLabel("Copy-Paste 概率:"),12,0); aug_layout.addWidget(self.copy_paste_spin,12,1)
aug_layout.addWidget(QLabel("关闭 Mosaic 轮数:"),13,0); aug_layout.addWidget(self.close_mosaic_spin,13,1)

param_layout.addWidget(aug_group)

# 按钮行
btn_layout = QHBoxLayout()
self.start_btn = QPushButton("🚀 开始训练")
self.stop_btn = QPushButton("⏹️ 停止训练")
self.stop_btn.setEnabled(False)
self.start_btn.clicked.connect(self.start_training)
self.stop_btn.clicked.connect(self.stop_training)
self.choose_data_btn = QPushButton("📂 选择数据YAML")
self.choose_model_btn = QPushButton("📂 选择模型YAML")
self.choose_weights_btn = QPushButton("📂 选择预训练权重")
self.choose_data_btn.clicked.connect(lambda: self.select_file(self.data_edit, "YAML (*.yaml)"))
self.choose_model_btn.clicked.connect(lambda: self.select_file(self.model_yaml_edit, "YAML (*.yaml)"))
self.choose_weights_btn.clicked.connect(lambda: self.select_file(self.weights_edit, "PT (*.pt)"))
btn_layout.addWidget(self.start_btn)
btn_layout.addWidget(self.stop_btn)
btn_layout.addWidget(self.choose_data_btn)
btn_layout.addWidget(self.choose_model_btn)
btn_layout.addWidget(self.choose_weights_btn)
btn_layout.addStretch()
param_layout.addLayout(btn_layout)

param_scroll.setWidget(param_widget)
main_layout.addWidget(param_scroll)

# 进度条和日志
self.progress_bar = QProgressBar()
self.progress_bar.setVisible(False)
main_layout.addWidget(self.progress_bar)

self.log_text = QTextEdit()
self.log_text.setReadOnly(True)
self.log_text.setMaximumHeight(150)
main_layout.addWidget(QLabel("📄 训练日志:"))
main_layout.addWidget(self.log_text)

# Tab 容器
self.tab_widget = QTabWidget()
main_layout.addWidget(self.tab_widget)

# TensorBoard 标签页
self.tb_view = QWebEngineView()
self.tab_widget.addTab(self.tb_view, "📈 TensorBoard")

# 实时指标标签页
self.metrics_widget = QWidget()
metrics_layout = QVBoxLayout(self.metrics_widget)
self.metrics_label = QLabel("等待训练开始…")
self.metrics_label.setAlignment(Qt.AlignCenter)
self.metrics_label.setFont(QFont("微软雅黑", 12))
metrics_layout.addWidget(self.metrics_label)
self.tab_widget.addTab(self.metrics_widget, "📊 实时指标")

# 图片测试标签页
self.test_widget = QWidget()
test_layout = QVBoxLayout(self.test_widget)
test_control = QHBoxLayout()
self.test_btn = QPushButton("🎲 随机抽取两张图片测试")
self.test_btn.clicked.connect(self.run_image_test)
test_control.addWidget(self.test_btn)
self.manual_weight_btn = QPushButton("📁 手动选择权重测试")
self.manual_weight_btn.clicked.connect(self.manual_test)
test_control.addWidget(self.manual_weight_btn)
self.choose_image_dir_btn = QPushButton("📁 手动选择图片目录")
self.choose_image_dir_btn.clicked.connect(self.manual_select_image_dir)
test_control.addWidget(self.choose_image_dir_btn)
test_control.addStretch()
test_layout.addLayout(test_control)

img_splitter = QSplitter(Qt.Horizontal)
self.best_label = QLabel("best.pt 预测结果")
self.best_label.setStyleSheet("border: 1px solid #ccc; background: #fafafa;")
self.best_label.setAlignment(Qt.AlignCenter)
self.last_label = QLabel("last.pt 预测结果")
self.last_label.setStyleSheet("border: 1px solid #ccc; background: #fafafa;")
self.last_label.setAlignment(Qt.AlignCenter)
img_splitter.addWidget(self.best_label)
img_splitter.addWidget(self.last_label)
test_layout.addWidget(img_splitter)
self.tab_widget.addTab(self.test_widget, "🖼️ 模型测试")

# ========== 新增: 历史记录标签页 ==========
self.history_widget = HistoryWidget()
self.tab_widget.addTab(self.history_widget, "📜 历史记录 & 曲线")

# 添加模型管理与导出标签页
self.model_manager_widget = ModelManagerWidget()
self.tab_widget.addTab(self.model_manager_widget, "🔧 模型管理 & 导出")

# 定时刷新 TensorBoard 页面
self.tb_timer = QTimer()
self.tb_timer.timeout.connect(self.refresh_tensorboard)
self.tb_timer.start(5000)

# ———————— 辅助函数 ————————
def select_file(self, line_edit, filter_str):
path, _ = QFileDialog.getOpenFileName(self, "选择文件", "", filter_str)
if path:
line_edit.setText(path)

def init_tensorboard_server(self):
try:
log_dir = "runs"
os.makedirs(log_dir, exist_ok=True)
self.tb_process = subprocess.Popen(
["tensorboard", "–logdir", log_dir, "–port", "6006", "–bind_all"],
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
)
time.sleep(2)
self.tb_view.setUrl(QUrl(self.tb_url))
self.log_text.append(f"[INFO] TensorBoard 已启动，监听目录: {log_dir}")
except Exception as e:
self.log_text.append(f"[警告] TensorBoard 启动失败: {e}，您可手动运行: tensorboard –logdir runs")

def refresh_tensorboard(self):
if self.tb_view.url().isEmpty():
self.tb_view.setUrl(QUrl(self.tb_url))

def get_train_config(self):
config = {
'data': self.data_edit.text(),
'model_yaml': self.model_yaml_edit.text(),
'pretrained_weights': self.weights_edit.text(),
'epochs': self.epochs_spin.value(),
'batch': self.batch_spin.value(),
'imgsz': self.imgsz_spin.value(),
'device': self.device_edit.text(),
'workers': self.workers_spin.value(),
'optimizer': self.optimizer_combo.currentText(),
'project': 'runs/train',
'name': 'exp',
'pretrained': True,
'lr0': self.lr0_spin.value(),
'lrf': self.lrf_spin.value(),
'momentum': self.momentum_spin.value(),
'weight_decay': self.weight_decay_spin.value(),
'warmup_epochs': self.warmup_epochs_spin.value(),
'warmup_momentum': self.warmup_momentum_spin.value(),
'warmup_bias_lr': self.warmup_bias_lr_spin.value(),
'box_loss_gain': self.box_loss_spin.value(),
'cls_loss_gain': self.cls_loss_spin.value(),
'dfl_loss_gain': self.dfl_loss_spin.value(),
'hsv_h': self.hsv_h_spin.value(),
'hsv_s': self.hsv_s_spin.value(),
'hsv_v': self.hsv_v_spin.value(),
'degrees': self.degrees_spin.value(),
'translate': self.translate_spin.value(),
'scale': self.scale_spin.value(),
'shear': self.shear_spin.value(),
'perspective': self.perspective_spin.value(),
'flipud': self.flipud_spin.value(),
'fliplr': self.fliplr_spin.value(),
'mosaic': self.mosaic_spin.value(),
'mixup': self.mixup_spin.value(),
'copy_paste': self.copy_paste_spin.value(),
'close_mosaic': self.close_mosaic_spin.value(),
}
return config

def start_training(self):
config = self.get_train_config()
if not os.path.exists(config['data']):
QMessageBox.warning(self, "错误", "数据集YAML文件不存在")
return
if not os.path.exists(config['model_yaml']):
QMessageBox.warning(self, "错误", "模型YAML文件不存在")
return
self.start_btn.setEnabled(False)
self.stop_btn.setEnabled(True)
self.progress_bar.setVisible(True)
self.progress_bar.setValue(0)
self.log_text.clear()
self.train_worker = TrainWorker(config)
self.train_worker.log_signal.connect(self.append_log)
self.train_worker.progress_signal.connect(self.update_progress)
self.train_worker.metric_signal.connect(self.update_metrics_display)
self.train_worker.finished_signal.connect(self.training_finished)
self.train_worker.save_dir_signal.connect(self.set_save_dir)
self.train_worker.session_created_signal.connect(self.on_session_created)
self.train_worker.start()

def on_session_created(self, session_id):
self.append_log(f"[INFO] 数据库会话已创建，ID: {session_id}")

def set_save_dir(self, save_dir):
self.train_dir = save_dir
self.best_pt_path = os.path.join(self.train_dir, 'weights', 'best.pt')
self.last_pt_path = os.path.join(self.train_dir, 'weights', 'last.pt')
self.append_log(f"[INFO] 实际保存目录: {self.train_dir}")
if os.path.exists(self.best_pt_path):
self.append_log("[INFO] 已发现 best.pt")
if os.path.exists(self.last_pt_path):
self.append_log("[INFO] 已发现 last.pt")

def stop_training(self):
if self.train_worker and self.train_worker.isRunning():
self.train_worker.stop()
self.append_log("[INFO] 正在停止训练…")

def append_log(self, msg):
self.log_text.append(msg)
self.log_text.ensureCursorVisible()

def update_progress(self, current, total):
self.progress_bar.setMaximum(total)
self.progress_bar.setValue(current)

def update_metrics_display(self, metrics):
epoch = metrics.get('epoch', 0)
loss = to_scalar(metrics.get('loss', 0.0))
box = to_scalar(metrics.get('box_loss', 0.0))
cls = to_scalar(metrics.get('cls_loss', 0.0))
dfl = to_scalar(metrics.get('dfl_loss', 0.0))
text = f"Epoch {epoch} | 总损失: {loss:.4f} | Box: {box:.4f} | Cls: {cls:.4f} | DFL: {dfl:.4f}"
self.metrics_label.setText(text)

def training_finished(self, success, msg):
self.start_btn.setEnabled(True)
self.stop_btn.setEnabled(False)
self.progress_bar.setVisible(False)
self.append_log(msg)
if success:
QMessageBox.information(self, "训练完成", msg)
self.run_image_test()
# 刷新历史记录标签页的会话列表
self.history_widget.load_session_list()

# ———————— 图片测试功能（增强路径解析 + 手动选择目录） ————————
def find_latest_train_dir(self):
base_dirs = ["runs/train", "runs/detect/train"]
latest = None
latest_time = 0
for base in base_dirs:
if os.path.exists(base):
for exp in os.listdir(base):
exp_path = os.path.join(base, exp)
if os.path.isdir(exp_path) and exp.startswith("exp"):
mtime = os.path.getmtime(exp_path)
if mtime > latest_time:
latest_time = mtime
latest = exp_path
return latest

def manual_test(self):
file_path, _ = QFileDialog.getOpenFileName(self, "选择权重文件", "", "PT files (*.pt)")
if not file_path:
return
self.run_image_test(manual_best=file_path, manual_last=None)

def manual_select_image_dir(self):
dir_path = QFileDialog.getExistingDirectory(self, "选择训练图片目录")
if dir_path:
self.manual_image_dir = dir_path
self.append_log(f"[INFO] 已手动设置测试图片目录: {dir_path}")

def get_train_image_dir(self):
data_yaml = self.data_edit.text()
if not os.path.exists(data_yaml):
return None

try:
with open(data_yaml, 'r') as f:
data_cfg = yaml.safe_load(f)

train_path = data_cfg.get('train')
base_path = data_cfg.get('path', '')

candidates = []
if train_path:
if not os.path.isabs(train_path):
base_dir = os.path.dirname(os.path.abspath(data_yaml))
candidates.append(os.path.join(base_dir, train_path))
else:
candidates.append(train_path)
if base_path:
if not os.path.isabs(base_path):
base_dir = os.path.dirname(os.path.abspath(data_yaml))
base_abs = os.path.join(base_dir, base_path)
else:
base_abs = base_path
candidates.append(os.path.join(base_abs, train_path))

if not train_path:
base_dir = os.path.dirname(os.path.abspath(data_yaml))
candidates = [
os.path.join(base_dir, 'images', 'train'),
os.path.join(base_dir, 'train'),
os.path.join(base_dir, 'data', 'train'),
]

for cand in candidates:
if os.path.exists(cand):
return cand

return None
except Exception as e:
self.append_log(f"[错误] 解析 data.yaml 失败: {e}")
return None

def run_image_test(self, manual_best=None, manual_last=None):
if manual_best and os.path.exists(manual_best):
best_path = manual_best
last_path = None
else:
if not self.train_dir or not os.path.exists(self.train_dir):
latest = self.find_latest_train_dir()
if latest:
self.set_save_dir(latest)
if not self.best_pt_path or not os.path.exists(self.best_pt_path):
QMessageBox.warning(self, "测试失败", "未找到 best.pt，请先训练或使用「手动选择权重测试」")
return
best_path = self.best_pt_path
last_path = self.last_pt_path if self.last_pt_path and os.path.exists(self.last_pt_path) else None

train_img_dir = None
if self.manual_image_dir and os.path.exists(self.manual_image_dir):
train_img_dir = self.manual_image_dir
self.append_log(f"[INFO] 使用手动选择的图片目录: {train_img_dir}")
else:
train_img_dir = self.get_train_image_dir()
if not train_img_dir:
reply = QMessageBox.question(self, "提示", "无法自动定位训练图片目录，是否手动选择？",
QMessageBox.Yes | QMessageBox.No)
if reply == QMessageBox.Yes:
dir_path = QFileDialog.getExistingDirectory(self, "选择训练图片目录")
if dir_path:
train_img_dir = dir_path
self.manual_image_dir = dir_path
else:
return
else:
return
else:
self.append_log(f"[INFO] 自动定位训练图片目录: {train_img_dir}")

img_ext = ('.jpg','.jpeg','.png','.bmp')
all_imgs = [p for p in Path(train_img_dir).iterdir() if p.suffix.lower() in img_ext]
if len(all_imgs) < 2:
self.append_log(f"[错误] 训练图片不足2张，实际{len(all_imgs)}")
return
selected = random.sample(all_imgs, 2)

try:
model_best = YOLO(best_path)
model_last = YOLO(last_path) if last_path else None
except Exception as e:
self.append_log(f"[错误] 加载模型失败: {e}")
return

res_best = [model_best.predict(str(p), conf=0.25, save=False)[0] for p in selected]
res_last = [model_last.predict(str(p), conf=0.25, save=False)[0] for p in selected] if model_last else []

best_pix = self.combine_images(res_best, selected, target_width=640)
last_pix = self.combine_images(res_last, selected, target_width=640) if res_last else None
if best_pix:
scaled = best_pix.scaled(self.best_label.width(), self.best_label.height(), Qt.KeepAspectRatio)
self.best_label.setPixmap(scaled)
if last_pix:
scaled = last_pix.scaled(self.last_label.width(), self.last_label.height(), Qt.KeepAspectRatio)
self.last_label.setPixmap(scaled)
elif not last_path:
self.last_label.setText("未选择 last.pt")

def combine_images(self, results, img_paths, target_width=640):
"""
将多张预测结果垂直拼接，自动统一宽度

Args:
results: YOLO 预测结果列表
img_paths: 对应图片路径列表
target_width: 统一的目标宽度（像素），如果为 None 则自动使用所有图片中的最大宽度
"""
imgs = []
for res, path in zip(results, img_paths):
img = res.plot() # BGR numpy array
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 添加图片文件名
cv2.putText(img_rgb, os.path.basename(path), (10, 30),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
imgs.append(img_rgb)

if not imgs:
return None

# 计算目标宽度
if target_width is None:
# 使用所有图片中的最大宽度
target_width = max(img.shape[1] for img in imgs)

# 缩放所有图片到目标宽度（保持高宽比）
resized_imgs = []
for img in imgs:
h, w = img.shape[:2]
if w != target_width:
scale = target_width / w
new_h = int(h * scale)
# 使用 OpenCV 进行缩放
resized = cv2.resize(img, (target_width, new_h), interpolation=cv2.INTER_AREA)
resized_imgs.append(resized)
else:
resized_imgs.append(img)

# 垂直拼接
concat = np.vstack(resized_imgs)
h, w, ch = concat.shape
qimg = QImage(concat.data, w, h, ch * w, QImage.Format_RGB888)
return QPixmap.fromImage(qimg)

def closeEvent(self, event):
if self.train_worker and self.train_worker.isRunning():
self.train_worker.stop()
self.train_worker.wait(2000)
if self.tb_process:
self.tb_process.terminate()
event.accept()

if __name__ == '__main__':
from PyQt5 import QtCore
QtCore.QCoreApplication.setAttribute(QtCore.Qt.AA_ShareOpenGLContexts)
app = QApplication(sys.argv)
app.setFont(QFont("微软雅黑", 9))
window = YOLOTrainGUI()
window.show()
sys.exit(app.exec_())

作者寄语：搭建工具平台的过程本身就是深度学习能力的重要体现——它要求你不仅理解算法原理，还要具备系统工程思维。希望本文能为你构建自己的AI工具链提供有价值的参考。如有任何问题或建议，欢迎在评论区交流讨论。如果需要我的完整工程源码，可后台私信（发送文章，名称或者链接可获取）。