以下是针对 分批次加载大数据集并分阶段训练 YOLOv5 模型 的完整代码示例和操作步骤：
1. 数据集分割脚本
将 20 万张图片和标签文件分割为 4 个子集（每个 5 万张），保持类别分布平衡。
import os
import shutil
import numpy as np
from sklearn.model_selection import train_test_splitdef split_dataset(data_dir, output_base, n_splits=4):"""分割数据集为多个子集Args:data_dir: 原始数据集目录，包含 images 和 labels 子目录output_base: 输出根目录（自动生成 train1, train2...）n_splits: 分割数量"""# 获取所有图片文件名（不带扩展名）image_files = [f.split('.')[0] for f in os.listdir(f"{data_dir}/images")]image_files = sorted(image_files)# 分层随机分割（假设每个图片的标签文件存在）splits = np.array_split(image_files, n_splits)# 创建子目录并复制文件for i, split in enumerate(splits, 1):split_dir = f"{output_base}/train{i}"os.makedirs(f"{split_dir}/images", exist_ok=True)os.makedirs(f"{split_dir}/labels", exist_ok=True)for name in split:# 复制图片src_img = f"{data_dir}/images/{name}.jpg"dst_img = f"{split_dir}/images/{name}.jpg"shutil.copy(src_img, dst_img)# 复制标签src_label = f"{data_dir}/labels/{name}.txt"dst_label = f"{split_dir}/labels/{name}.txt"shutil.copy(src_label, dst_label)print(f"Subset {i} created: {len(split)} samples")# 使用示例
split_dataset(data_dir="datasets/original",output_base="datasets/split",n_splits=4
)
2. 分阶段训练脚本
依次在每个子集上训练模型，并继承前一阶段的权重。
import os
import yaml# 基础配置
base_cfg = {"path": "datasets/split","train": "train1/images", # 初始训练集路径（将被动态修改）"val": "val/images", # 验证集路径（保持不变）"names": ["plate"] # 类别名称
}# 保存动态 data.yaml
def update_data_yaml(stage):base_cfg["train"] = f"train{stage}/images"with open("data.yaml", 'w') as f:yaml.dump(base_cfg, f)# 分阶段训练
for stage in range(1, 5):print(f"\n=== Training Stage {stage} ===")# 1. 更新 data.yaml 中的训练集路径update_data_yaml(stage)# 2. 设置训练命令参数weights = "yolov5n.pt" if stage == 1 else f"runs/stage{stage-1}/weights/last.pt"cmd = f"""python train.py \--weights {weights} \--data data.yaml \--img 320 \--batch 8 \--epochs 50 \--project runs/stage{stage} \--name exp \--exist-ok"""# 3. 执行训练os.system(cmd)
3. 关键说明
(1) 数据集目录结构
datasets/split/
├── train1/
│ ├── images/ # 5万张图片
│ └── labels/ # 对应标签
├── train2/
├── train3/
├── train4/
└── val/ # 独立的验证集（需预先准备）
(2) 训练流程
- 阶段1：在
train1上训练，从yolov5n.pt开始。 - 阶段2：在
train2上训练，加载runs/stage1/weights/last.pt。 - 阶段3：在
train3上训练，加载runs/stage2/weights/last.pt。 - 阶段4：在
train4上训练，加载runs/stage3/weights/last.pt。
(3) 参数优化
- 批次大小：
--batch 8（根据显存调整）。 - 分辨率：
--img 320（更低分辨率节省显存）。 - 训练周期：
--epochs 50（每个子集训练50轮）。
4. 注意事项
-
验证集一致性：
所有阶段使用相同的验证集（val目录），确保评估指标可比性。 -
学习率调整：
在后续阶段可降低学习率（如--lr0 0.001），避免破坏已有特征。 -
权重选择：
每个阶段结束后，检查验证集 mAP，选择最佳权重（best.pt）而非last.pt。 -
资源监控：
使用nvidia-smi和free -h监控显存和内存使用，必要时减少--workers。
5. 扩展方案
若需进一步优化：
- 增量训练：在完整数据集上微调最终权重。
python train.py --weights runs/stage4/weights/best.pt --data data_full.yaml ... - 分布式训练：使用多 GPU 加速。
python -m torch.distributed.launch --nproc_per_node 2 train.py ...