YOLO训练策略与优化技巧：从数据准备到模型收敛

import os
import json
import cv2

# COCO格式转YOLO格式
def coco_to_yolo(coco_json_path, images_dir, output_dir):
    # 创建输出目录
    os.makedirs(os.path.join(output_dir, 'images'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels'), exist_ok=True)
    
    # 加载COCO JSON文件
    with open(coco_json_path, 'r') as f:
        coco_data = json.load(f)
    
    # 获取类别映射
    categories = {cat['id']: cat['name'] for cat in coco_data['categories']}
    class_mapping = {cat_id: i for i, cat_id in enumerate(sorted(categories.keys()))}
    
    # 转换标注
    for img in coco_data['images']:
        img_id = img['id']
        img_name = img['file_name']
        img_width = img['width']
        img_height = img['height']
        
        # 复制图像文件
        src_img_path = os.path.join(images_dir, img_name)
        dst_img_path = os.path.join(output_dir, 'images', img_name)
        if os.path.exists(src_img_path):
            os.system(f'cp {src_img_path} {dst_img_path}')
        
        # 生成标注文件
        label_path = os.path.join(output_dir, 'labels', os.path.splitext(img_name)[0] + '.txt')
        with open(label_path, 'w') as f:
            # 查找该图像的所有标注
            for ann in coco_data['annotations']:
                if ann['image_id'] == img_id:
                    # 获取类别ID
                    cat_id = ann['category_id']
                    if cat_id not in class_mapping:
                        continue
                    class_id = class_mapping[cat_id]
                    
                    # 获取边界框
                    bbox = ann['bbox']  # [x, y, width, height]
                    x, y, w, h = bbox
                    
                    # 转换为YOLO格式：x_center, y_center, width, height（归一化）
                    x_center = (x + w / 2) / img_width
                    y_center = (y + h / 2) / img_height
                    norm_w = w / img_width
                    norm_h = h / img_height
                    
                    # 写入标注文件
                    f.write(f'{class_id} {x_center:.6f} {y_center:.6f} {norm_w:.6f} {norm_h:.6f}\n')
    
    # 生成data.yaml文件
    data_yaml = {
        'train': os.path.join(output_dir, 'images'),
        'val': os.path.join(output_dir, 'images'),  # 实际应分开
        'nc': len(class_mapping),
        'names': [categories[cat_id] for cat_id in sorted(class_mapping.keys())]
    }
    
    with open(os.path.join(output_dir, 'data.yaml'), 'w') as f:
        yaml.dump(data_yaml, f, default_flow_style=False)
    
    print(f'转换完成，共处理{len(coco_data["images"])}张图像')
    return data_yaml

import os
import cv2
import numpy as np

def clean_dataset(images_dir, labels_dir, output_dir, min_size=10):
    # 创建输出目录
    os.makedirs(os.path.join(output_dir, 'images'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels'), exist_ok=True)
    
    cleaned_count = 0
    
    # 遍历所有图像文件
    for img_name in os.listdir(images_dir):
        if img_name.endswith(('.jpg', '.jpeg', '.png')):
            img_path = os.path.join(images_dir, img_name)
            label_path = os.path.join(labels_dir, os.path.splitext(img_name)[0] + '.txt')
            
            # 检查图像和标注是否存在
            if not os.path.exists(img_path) or not os.path.exists(label_path):
                continue
            
            # 读取图像
            img = cv2.imread(img_path)
            if img is None:
                continue
            
            img_height, img_width = img.shape[:2]
            
            # 读取标注
            with open(label_path, 'r') as f:
                lines = f.readlines()
            
            cleaned_lines = []
            for line in lines:
                line = line.strip()
                if not line:
                    continue
                
                parts = line.split()
                if len(parts) != 5:
                    continue
                
                class_id, x_center, y_center, width, height = map(float, parts)
                
                # 检查坐标是否合法
                if not (0 <= x_center <= 1 and 0 <= y_center <= 1 and 0 <= width <= 1 and 0 <= height <= 1):
                    continue
                
                # 检查目标大小
                abs_width = width * img_width
                abs_height = height * img_height
                if abs_width < min_size or abs_height < min_size:
                    continue
                
                cleaned_lines.append(line)
            
            # 如果有有效标注，保存清洗后的图像和标注
            if cleaned_lines:
                # 保存图像
                dst_img_path = os.path.join(output_dir, 'images', img_name)
                cv2.imwrite(dst_img_path, img)
                
                # 保存标注
                dst_label_path = os.path.join(output_dir, 'labels', os.path.splitext(img_name)[0] + '.txt')
                with open(dst_label_path, 'w') as f:
                    f.write('\n'.join(cleaned_lines) + '\n')
                
                cleaned_count += 1
    
    print(f'数据清洗完成，共保留{cleaned_count}张有效图像')

import numpy as np
import cv2
import random

def mosaic_augmentation(images, labels, img_size=640):
    # 创建输出图像
    output_img = np.zeros((img_size, img_size, 3), dtype=np.uint8)
    output_labels = []
    
    # 随机生成拼接中心点
    xc, yc = [int(random.uniform(img_size * 0.25, img_size * 0.75)) for _ in range(2)]
    
    # 遍历4张图像
    for i in range(4):
        img, label = images[i], labels[i]
        
        # 调整图像大小
        h, w = img.shape[:2]
        scale = min(img_size / h, img_size / w)
        nh, nw = int(h * scale), int(w * scale)
        resized_img = cv2.resize(img, (nw, nh))
        
        # 确定图像在输出中的位置
        if i == 0:  # 左上
            x1a, y1a, x2a, y2a = 0, 0, xc, yc
            x1b, y1b, x2b, y2b = 0, 0, nw, nh
        elif i == 1:  # 右上
            x1a, y1a, x2a, y2a = xc, 0, img_size, yc
            x1b, y1b, x2b, y2b = nw - (x2a - x1a), 0, nw, nh
        elif i == 2:  # 左下
            x1a, y1a, x2a, y2a = 0, yc, xc, img_size
            x1b, y1b, x2b, y2b = 0, nh - (y2a - y1a), nw, nh
        else:  # 右下
            x1a, y1a, x2a, y2a = xc, yc, img_size, img_size
            x1b, y1b, x2b, y2b = nw - (x2a - x1a), nh - (y2a - y1a), nw, nh
        
        # 裁剪并复制图像到输出
        output_img[y1a:y2a, x1a:x2a] = resized_img[y1b:y2b, x1b:x2b]
        
        # 调整标注
        if len(label) > 0:
            # 计算坐标转换比例
            sx = (x2a - x1a) / (x2b - x1b)
            sy = (y2a - y1a) / (y2b - y1b)
            
            # 调整边界框坐标
            for l in label:
                class_id, x_center, y_center, width, height = l
                
                # 转换为绝对坐标
                x_center_abs = x_center * nw
                y_center_abs = y_center * nh
                width_abs = width * nw
                height_abs = height * nh
                
                # 调整坐标到输出图像
                x_center_abs = x_center_abs - x1b
                y_center_abs = y_center_abs - y1b
                
                # 检查目标是否在拼接区域内
                if x_center_abs < 0 or x_center_abs > (x2b - x1b) or y_center_abs < 0 or y_center_abs > (y2b - y1b):
                    continue
                
                # 转换到输出图像坐标
                x_center_abs = x_center_abs * sx + x1a
                y_center_abs = y_center_abs * sy + y1a
                width_abs = width_abs * sx
                height_abs = height_abs * sy
                
                # 转换回归一化坐标
                x_center = x_center_abs / img_size
                y_center = y_center_abs / img_size
                width = width_abs / img_size
                height = height_abs / img_size
                
                # 检查坐标是否合法
                if 0 <= x_center <= 1 and 0 <= y_center <= 1 and 0 <= width <= 1 and 0 <= height <= 1:
                    output_labels.append([class_id, x_center, y_center, width, height])
    
    return output_img, output_labels

def mixup_augmentation(img1, label1, img2, label2, alpha=1.0):
    # 生成混合比例
    lam = np.random.beta(alpha, alpha)
    
    # 混合图像
    mixed_img = lam * img1 + (1 - lam) * img2
    mixed_img = mixed_img.astype(np.uint8)
    
    # 混合标注
    mixed_labels = []
    for l in label1:
        mixed_labels.append(l + [lam])
    for l in label2:
        mixed_labels.append(l + [1 - lam])
    
    return mixed_img, mixed_labels

def cutmix_augmentation(img1, label1, img2, label2, alpha=1.0):
    # 生成混合比例
    lam = np.random.beta(alpha, alpha)
    
    # 获取图像尺寸
    h, w = img1.shape[:2]
    
    # 生成裁剪区域
    cut_rat = np.sqrt(1.0 - lam)
    cut_w = int(w * cut_rat)
    cut_h = int(h * cut_rat)
    
    # 随机生成裁剪区域的中心点
    cx = np.random.randint(w)
    cy = np.random.randint(h)
    
    # 计算裁剪区域的边界
    x1 = np.clip(cx - cut_w // 2, 0, w)
    y1 = np.clip(cy - cut_h // 2, 0, h)
    x2 = np.clip(cx + cut_w // 2, 0, w)
    y2 = np.clip(cy + cut_h // 2, 0, h)
    
    # 执行CutMix
    mixed_img = img1.copy()
    mixed_img[y1:y2, x1:x2] = img2[y1:y2, x1:x2]
    
    # 调整混合比例
    lam = 1 - ((x2 - x1) * (y2 - y1) / (w * h))
    
    # 混合标注
    mixed_labels = []
    for l in label1:
        mixed_labels.append(l + [lam])
    for l in label2:
        mixed_labels.append(l + [1 - lam])
    
    return mixed_img, mixed_labels

class AdaptiveAugmentation:
    def __init__(self, augmentations, initial_strength=0.5, max_strength=1.0, min_strength=0.0):
        self.augmentations = augmentations  # 数据增强方法列表
        self.strength = initial_strength  # 当前增强强度
        self.max_strength = max_strength  # 最大增强强度
        self.min_strength = min_strength  # 最小增强强度
        self.last_loss = None  # 上一轮损失
    
    def update_strength(self, current_loss):
        # 根据损失变化调整增强强度
        if self.last_loss is None:
            self.last_loss = current_loss
            return
        
        # 如果损失下降，适当增加增强强度
        if current_loss < self.last_loss:
            self.strength = min(self.max_strength, self.strength + 0.05)
        # 如果损失上升，适当降低增强强度
        else:
            self.strength = max(self.min_strength, self.strength - 0.1)
        
        self.last_loss = current_loss
    
    def apply(self, img, label):
        # 根据当前强度选择增强方法
        for aug in self.augmentations:
            if np.random.rand() < self.strength:
                img, label = aug(img, label)
        
        return img, label

def iou_loss(pred_boxes, target_boxes):
    # 计算IoU
    iou = box_iou(pred_boxes, target_boxes)
    # IoU损失
    loss = 1 - iou
    return loss.mean()

# 计算IoU
def box_iou(box1, box2):
    # box1: [N, 4] (x1, y1, x2, y2)
    # box2: [M, 4] (x1, y1, x2, y2)
    
    # 计算交集区域
    lt = torch.max(box1[:, None, :2], box2[:, :2])  # [N, M, 2]
    rb = torch.min(box1[:, None, 2:], box2[:, 2:])  # [N, M, 2]
    
    wh = (rb - lt).clamp(min=0)  # [N, M, 2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N, M]
    
    # 计算面积
    area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])  # [N]
    area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])  # [M]
    
    # 计算并集
    union = area1[:, None] + area2 - inter  # [N, M]
    
    # 计算IoU
    iou = inter / (union + 1e-16)  # [N, M]
    
    return iou

def ciou_loss(pred_boxes, target_boxes):
    # 计算IoU
    iou = box_iou(pred_boxes, target_boxes)
    
    # 计算中心点距离
    pred_center = (pred_boxes[:, :2] + pred_boxes[:, 2:]) / 2
    target_center = (target_boxes[:, :2] + target_boxes[:, 2:]) / 2
    center_distance = torch.sum((pred_center - target_center) ** 2, dim=-1)
    
    # 计算最小外接矩形的对角线长度
    lt = torch.min(pred_boxes[:, :2], target_boxes[:, :2])
    rb = torch.max(pred_boxes[:, 2:], target_boxes[:, 2:])
    diagonal_distance = torch.sum((rb - lt) ** 2, dim=-1)
    
    # 计算宽高比
    pred_w = pred_boxes[:, 2] - pred_boxes[:, 0]
    pred_h = pred_boxes[:, 3] - pred_boxes[:, 1]
    target_w = target_boxes[:, 2] - target_boxes[:, 0]
    target_h = target_boxes[:, 3] - target_boxes[:, 1]
    
    v = (4 / (torch.pi ** 2)) * torch.pow(torch.atan(target_w / (target_h + 1e-16)) - torch.atan(pred_w / (pred_h + 1e-16)), 2)
    
    # 计算权重
    alpha = v / (1 - iou + v + 1e-16)
    
    # CIoU损失
    ciou = iou - (center_distance / diagonal_distance + alpha * v)
    loss = 1 - ciou
    
    return loss.mean()

def focal_loss(pred_cls, target_cls, alpha=0.25, gamma=2.0):
    # pred_cls: [B, N, C] 预测的类别概率
    # target_cls: [B, N, C] 真实类别标签（one-hot编码）
    
    # 计算交叉熵损失
    ce_loss = torch.nn.functional.cross_entropy(pred_cls, target_cls, reduction='none')
    
    # 计算类别概率
    p = torch.exp(-ce_loss)
    
    # 计算Focal Loss
    focal_loss = alpha * (1 - p) ** gamma * ce_loss
    
    return focal_loss.mean()

class DynamicLoss(nn.Module):
    def __init__(self, initial_box_weight=5.0, initial_cls_weight=1.0, initial_obj_weight=1.0):
        super(DynamicLoss, self).__init__()
        self.box_weight = initial_box_weight
        self.cls_weight = initial_cls_weight
        self.obj_weight = initial_obj_weight
        
    def forward(self, pred, target, epoch, max_epoch):
        # 动态调整损失权重
        # 随着训练进行，逐渐增加边界框损失的权重
        box_weight = self.box_weight * (epoch / max_epoch) * 2
        
        # 计算各部分损失
        box_loss = ciou_loss(pred['boxes'], target['boxes'])
        cls_loss = focal_loss(pred['cls'], target['cls'])
        obj_loss = binary_cross_entropy(pred['obj'], target['obj'])
        
        # 总损失
        total_loss = box_weight * box_loss + self.cls_weight * cls_loss + self.obj_weight * obj_loss
        
        return total_loss, {
            'box_loss': box_loss.item(),
            'cls_loss': cls_loss.item(),
            'obj_loss': obj_loss.item(),
            'total_loss': total_loss.item()
        }

class SGD:
    def __init__(self, params, lr=0.01, momentum=0, weight_decay=0):
        self.params = params
        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay
        self.velocities = [torch.zeros_like(p) for p in params]
    
    def step(self, grads):
        for i, (p, g, v) in enumerate(zip(self.params, grads, self.velocities)):
            # 权重衰减
            if self.weight_decay > 0:
                g = g + self.weight_decay * p
            
            # 动量更新
            v = self.momentum * v + g
            self.velocities[i] = v
            
            # 参数更新
            p.data -= self.lr * v

class Adam:
    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
        self.params = params
        self.lr = lr
        self.beta1 = betas[0]
        self.beta2 = betas[1]
        self.eps = eps
        self.weight_decay = weight_decay
        self.t = 0
        self.m = [torch.zeros_like(p) for p in params]
        self.v = [torch.zeros_like(p) for p in params]
    
    def step(self, grads):
        self.t += 1
        
        for i, (p, g, m, v) in enumerate(zip(self.params, grads, self.m, self.v)):
            # 权重衰减
            if self.weight_decay > 0:
                g = g + self.weight_decay * p
            
            # 一阶矩估计
            m = self.beta1 * m + (1 - self.beta1) * g
            
            # 二阶矩估计
            v = self.beta2 * v + (1 - self.beta2) * g ** 2
            
            # 偏差修正
            m_hat = m / (1 - self.beta1 ** self.t)
            v_hat = v / (1 - self.beta2 ** self.t)
            
            # 参数更新
            p.data -= self.lr * m_hat / (torch.sqrt(v_hat) + self.eps)
            
            # 更新动量和速度
            self.m[i] = m
            self.v[i] = v

def cosine_annealing_lr(initial_lr, current_epoch, max_epoch, min_lr=0):
    # 余弦退火学习率
    lr = min_lr + 0.5 * (initial_lr - min_lr) * (1 + torch.cos(torch.tensor(current_epoch / max_epoch * torch.pi)))
    return lr

def warmup_lr(initial_lr, current_epoch, warmup_epochs):
    # Warmup学习率
    if current_epoch < warmup_epochs:
        lr = initial_lr * (current_epoch / warmup_epochs)
    else:
        lr = initial_lr
    return lr

def combined_lr(initial_lr, current_epoch, max_epoch, warmup_epochs, min_lr=0):
    # 先使用Warmup，再使用余弦退火
    if current_epoch < warmup_epochs:
        lr = initial_lr * (current_epoch / warmup_epochs)
    else:
        # 余弦退火
        epoch = current_epoch - warmup_epochs
        max_epoch = max_epoch - warmup_epochs
        lr = min_lr + 0.5 * (initial_lr - min_lr) * (1 + torch.cos(torch.tensor(epoch / max_epoch * torch.pi)))
    return lr

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from ultralytics import YOLO
import yaml

# 加载数据集配置
def load_data_config(data_yaml):
    with open(data_yaml, 'r') as f:
        data_config = yaml.safe_load(f)
    return data_config

# 训练函数
def train_yolo(data_yaml, epochs=100, batch_size=16, lr0=1e-3, model_name='yolov10n.pt'):
    # 加载数据集配置
    data_config = load_data_config(data_yaml)
    
    # 加载预训练模型
    model = YOLO(model_name)
    
    # 训练配置
    train_config = {
        'data': data_yaml,
        'epochs': epochs,
        'batch': batch_size,
        'lr0': lr0,
        'lrf': 1e-4,
        'weight_decay': 5e-4,
        'warmup_epochs': 3,
        'warmup_momentum': 0.8,
        'box': 7.5,
        'cls': 0.5,
        'dfl': 1.5,
        'device': 'cuda' if torch.cuda.is_available() else 'cpu',
        'workers': 8,
        'project': 'runs/train',
        'name': 'exp',
        'exist_ok': True,
        'verbose': True
    }
    
    # 开始训练
    results = model.train(**train_config)
    
    # 验证模型
    metrics = model.val()
    print(f"mAP@0.5: {metrics.box.map50:.2f}")
    print(f"mAP@0.5-0.95: {metrics.box.map:.2f}")
    
    # 导出模型
    model.export(format='onnx')
    
    return results, metrics

# 运行训练
if __name__ == "__main__":
    data_yaml = 'path/to/data.yaml'
    results, metrics = train_yolo(data_yaml, epochs=100, batch_size=16, lr0=1e-3)

from torch.utils.tensorboard import SummaryWriter

# 创建TensorBoard写入器
writer = SummaryWriter('runs/yolo_train')

# 在训练循环中记录损失
def train_loop(dataloader, model, loss_fn, optimizer, epoch, writer):
    size = len(dataloader.dataset)
    model.train()
    
    for batch, (X, y) in enumerate(dataloader):
        # 前向传播
        pred = model(X)
        loss = loss_fn(pred, y)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # 记录损失
        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
            writer.add_scalar('Training Loss', loss, epoch * len(dataloader) + batch)
    
# 在验证循环中记录性能
def val_loop(dataloader, model, loss_fn, epoch, writer):
    model.eval()
    val_loss = 0
    correct = 0
    
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            val_loss += loss_fn(pred, y).item()
            # 计算准确率
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    
    val_loss /= len(dataloader)
    correct /= len(dataloader.dataset)
    
    print(f"Validation Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {val_loss:>8f} \n")
    writer.add_scalar('Validation Loss', val_loss, epoch)
    writer.add_scalar('Validation Accuracy', correct, epoch)

epoch: 1/100, train/box_loss: 0.0523, train/cls_loss: 0.3452, train/dfl_loss: 0.2345, train/loss: 0.6320, train/Instances: 123.45, val/box_loss: 0.0489, val/cls_loss: 0.3123, val/dfl_loss: 0.2109, val/loss: 0.5721, val/Instances: 118.76, lr/pg0: 0.000100, lr/pg1: 0.000100, lr/pg2: 0.000100
epoch: 2/100, train/box_loss: 0.0476, train/cls_loss: 0.3123, train/dfl_loss: 0.2012, train/loss: 0.5611, train/Instances: 125.67, val/box_loss: 0.0456, val/cls_loss: 0.2890, val/dfl_loss: 0.1987, val/loss: 0.5333, val/Instances: 120.34, lr/pg0: 0.000200, lr/pg1: 0.000200, lr/pg2: 0.000200
...
epoch: 100/100, train/box_loss: 0.0123, train/cls_loss: 0.0876, train/dfl_loss: 0.0543, train/loss: 0.1542, train/Instances: 130.23, val/box_loss: 0.0134, val/cls_loss: 0.0987, val/dfl_loss: 0.0654, val/loss: 0.1775, val/Instances: 125.67, lr/pg0: 0.000010, lr/pg1: 0.000010, lr/pg2: 0.000010

# 基础训练
python train.py --data data.yaml --epochs 100 --batch 16 --model yolov10n.pt

# 使用混合精度训练
python train.py --data data.yaml --epochs 100 --batch 16 --model yolov10n.pt --amp

# 使用多GPU训练
python -m torch.distributed.run --nproc_per_node 4 train.py --data data.yaml --epochs 100 --batch 64 --model yolov10n.pt --distributed

# 恢复训练
python train.py --data data.yaml --epochs 100 --batch 16 --model runs/train/exp/weights/last.pt --resume

# 创建虚拟环境
conda create -n yolo python=3.9
conda activate yolo

# 安装依赖
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install opencv-python numpy matplotlib pandas
pip install ultralytics  # YOLOv8/9/10官方库
pip install timm  # 主流骨干网络库
pip install onnx onnxruntime  # ONNX支持
pip install tensorboard  # 可视化工具
pip install pycocotools  # COCO数据集支持