184 lines
6.0 KiB
Python
184 lines
6.0 KiB
Python
|
||
"""
|
||
YOLOv12 绝缘子检测优化训练脚本（最终稳定版）
|
||
主要改进：
|
||
1. 适配YOLOv12 参数规范
|
||
2. 移除已弃用参数
|
||
3. 优化训练配置
|
||
"""
|
||
|
||
import torch
|
||
from ultralytics import YOLO
|
||
import shutil
|
||
import os
|
||
import glob
|
||
from PIL import Image, ImageDraw
|
||
import numpy as np
|
||
|
||
# +++ 有效参数配置区块 +++
|
||
TRAIN_CONFIG = {
|
||
"stage1": {
|
||
"freeze": [0, 1, 2, 3, 4], # 冻结前5层
|
||
"epochs": 150, # 优化训练轮数
|
||
"batch": 48, # 增加批量大小
|
||
"lr0": 0.0005, # 优化初始学习率
|
||
"lrf": 0.1, # 提高最终学习率
|
||
"hsv_h": 0.01, # 色相增强幅度
|
||
"degrees": 10.0, # 增加旋转角度范围
|
||
"perspective": 0.005, # 增加透视变换
|
||
"flipud": 0.2, # 增加上下翻转概率
|
||
"optimizer": "SGD", # 优化器类型
|
||
"weight_decay": 0.001, # 权重衰减
|
||
"label_smoothing": 0.1, # 标签平滑
|
||
# "mosaic": True, # 启用Mosaic增强
|
||
"amp": True # 启用混合精度训练
|
||
},
|
||
"stage2": {
|
||
"epochs": 150, # 优化训练轮数
|
||
"batch": 48, # 增加批量大小
|
||
"lr0": 0.00005, # 优化初始学习率
|
||
"lrf": 0.01, # 提高最终学习率
|
||
"mixup": 0.2, # 增加MixUp增强系数
|
||
"close_mosaic": 10, # 最后10epoch关闭Mosaic
|
||
"optimizer": "AdamW", # 切换优化器为AdamW
|
||
"warmup_epochs": 10, # 增大学习率预热轮数
|
||
"amp": True # 启用混合精度训练
|
||
}
|
||
}
|
||
|
||
|
||
def validate_bbox_coordinates(label_path):
|
||
"""验证标注文件边界框坐标是否在标准化范围（0-1）"""
|
||
try:
|
||
with open(label_path) as f:
|
||
for line in f:
|
||
parts = list(map(float, line.strip().split()))
|
||
if len(parts) > 1:
|
||
x_center, y_center, width, height = parts[1:5]
|
||
assert 0 <= x_center <= 1 and 0 <= y_center <= 1, f"坐标超出范围：{label_path}"
|
||
assert 0 < width <= 1 and 0 < height <= 1, f"尺寸异常：{label_path}"
|
||
except Exception as e:
|
||
print(f"验证失败 {label_path}: {str(e)}")
|
||
|
||
|
||
def enhanced_data_check():
|
||
"""增强型数据验证（稳定版）"""
|
||
print("\n=== 执行增强数据验证 ===")
|
||
label_files = glob.glob("datasets/train/labels/*.txt")
|
||
class_dist = {}
|
||
|
||
# 类别分布分析
|
||
for lbl in label_files:
|
||
with open(lbl) as f:
|
||
for line in f:
|
||
class_id = int(line.strip().split()[0])
|
||
class_dist[class_id] = class_dist.get(class_id, 0) + 1
|
||
|
||
print(f"类别分布：{class_dist}")
|
||
|
||
# 坐标验证
|
||
for lbl in label_files:
|
||
validate_bbox_coordinates(lbl)
|
||
|
||
# 样本可视化
|
||
sample_count = 3
|
||
for i in range(sample_count):
|
||
img_path = f"datasets/train/images/{i:04d}.jpg"
|
||
lbl_path = f"datasets/train/labels/{i:04d}.txt"
|
||
try:
|
||
img = Image.open(img_path)
|
||
draw = ImageDraw.Draw(img)
|
||
with open(lbl_path) as f:
|
||
for line in f:
|
||
parts = list(map(float, line.strip().split()))
|
||
if len(parts) > 1:
|
||
x_center, y_center, width, height = parts[1:5]
|
||
# 转换为绝对坐标
|
||
w, h = img.size
|
||
x1 = (x_center - width/2) * w
|
||
y1 = (y_center - height/2) * h
|
||
x2 = (x_center + width/2) * w
|
||
y2 = (y_center + height/2) * h
|
||
draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
|
||
img.save(f"sample_{i}.jpg")
|
||
print(f"样本可视化已保存：sample_{i}.jpg")
|
||
except Exception as e:
|
||
print(f"可视化失败 {i}: {str(e)}")
|
||
|
||
def main():
|
||
# 初始化模型（正确方式）
|
||
model = YOLO("yolo12s.yaml") # 确保yaml文件存在
|
||
|
||
# === 阶段一：冻结训练 ===
|
||
print("\n=== 阶段一：特征提取层训练 ===")
|
||
stage1_results = model.train(
|
||
data="data.yaml",
|
||
**TRAIN_CONFIG["stage1"],
|
||
device=0,
|
||
workers=16,
|
||
box=5.0, # box损失权重
|
||
cls=1.0, # 分类损失权重
|
||
save_period=5, # 每5epoch保存检查点
|
||
patience=20, # 早停等待周期
|
||
deterministic=False # 禁用确定性模式
|
||
)
|
||
|
||
# === 阶段二：全网络微调 ===
|
||
print("\n=== 阶段二：全网络微调 ===")
|
||
stage2_results = model.train(
|
||
data="data.yaml",
|
||
**TRAIN_CONFIG["stage2"],
|
||
resume=True, # 从上一阶段继续
|
||
name="final_model", # 实验名称
|
||
shear=1.5, # 剪切变换幅度
|
||
copy_paste=0.05, # 复制粘贴增强
|
||
erasing=0.1, # 随机擦除概率
|
||
overlap_mask=False # 禁用掩码重叠
|
||
)
|
||
|
||
# === 模型验证 ===
|
||
print("\n=== 最终模型验证 ===")
|
||
metrics = model.val(
|
||
data="data.yaml",
|
||
conf=0.35, # 置信度阈值
|
||
iou=0.65, # IoU阈值
|
||
plots=True, # 生成评估图表
|
||
half=False # 禁用半精度验证
|
||
)
|
||
|
||
# === 模型导出 ===
|
||
print("\n=== 模型导出 ===")
|
||
|
||
# 导出为PT格式（TorchScript）
|
||
model.export(
|
||
format="torchscript", # PT格式
|
||
# dynamic=True,
|
||
simplify=True
|
||
)
|
||
|
||
# 导出为ONNX格式
|
||
model.export(
|
||
format="onnx", # ONNX格式
|
||
dynamic=True,
|
||
simplify=True
|
||
)
|
||
|
||
# 导出为TFLite格式
|
||
model.export(
|
||
format="tflite", # TFLite格式
|
||
# dynamic=True,
|
||
simplify=True
|
||
)
|
||
|
||
if __name__ == "__main__":
|
||
# 硬件优化配置
|
||
torch.set_float32_matmul_precision('high')
|
||
torch.backends.cudnn.benchmark = True
|
||
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
|
||
|
||
# 数据检查
|
||
enhanced_data_check()
|
||
|
||
# 启动训练
|
||
main()
|