当前位置：首页 > news >正文

Yolo迁移训练-带训练源码

news 来源：原创 2025/5/9 7:11:08

下载Git

拉下yolo模型

下载labelimg

准备训练集

迁移训练

继续训练

下载Git

Git - Downloading Package

拉下yolo模型

然后用克隆脚本拉下yolo模型

python clone_yolo.py

import os
import subprocess
import sys
import shutildef check_git_installed():"""检查Git是否已安装"""try:result = subprocess.run(['git', '--version'], capture_output=True, text=True)if result.returncode == 0:print(f"Git已安装: {result.stdout.strip()}")return Truereturn Falseexcept:return Falsedef clone_yolo():"""克隆YOLOv5仓库"""# 获取当前工作目录work_dir = os.path.dirname(os.path.abspath(__file__))# 检查是否已存在yolov5目录yolo_dir = os.path.join(work_dir, 'yolov5')if (os.path.exists(yolo_dir)):print(f"YOLO目录已存在({yolo_dir})。自动删除并重新克隆...")shutil.rmtree(yolo_dir, ignore_errors=True)# 克隆YOLOv5仓库 - 使用实时输出而不是捕获输出print("正在克隆YOLOv5仓库，这可能需要几分钟时间...")print("(下载进度将实时显示，请耐心等待)")try:# 使用实时输出方式运行git cloneresult = subprocess.run(['git', 'clone', 'https://github.com/ultralytics/yolov5.git'],cwd=work_dir,capture_output=False,  # 不捕获输出，让它直接显示在控制台text=True)if result.returncode == 0:print("YOLOv5仓库克隆成功!")# 安装依赖print("\n正在安装YOLOv5依赖...")requirements_file = os.path.join(yolo_dir, 'requirements.txt')if os.path.exists(requirements_file):try:subprocess.run([sys.executable, '-m', 'pip', 'install', '-r', requirements_file],check=True,capture_output=False  # 实时显示安装进度)print("YOLOv5依赖安装完成!")except subprocess.CalledProcessError as e:print(f"警告: 安装依赖时出错: {str(e)}")print("您可能需要手动安装依赖。")else:print(f"警告: 找不到requirements.txt文件({requirements_file})。")return Trueelse:print(f"克隆YOLOv5仓库失败，返回代码: {result.returncode}")return Falseexcept Exception as e:print(f"执行Git命令时出错: {e}")return Falseif __name__ == "__main__":if not check_git_installed():print("Git尚未安装或无法在PATH中找到。")print("请先运行 install_git.py 脚本安装Git，然后重新运行此脚本。")sys.exit(1)print("开始克隆YOLOv5仓库，请确保网络连接正常...")if clone_yolo():print("\nYOLOv5克隆和设置完成!")print("您现在可以使用 train_fruits.py 脚本进行水果识别的迁移学习。")else:print("\n克隆YOLO仓库失败，请检查错误信息并重试。")

下载labelimg

在下载yolo的过程中可以先去https://github.com/tzutalin/labelImg/releases下载给图像打标签的工具 labelimg

准备训练集

在 label 完后会生成 txt 文件（注意同一个水果图片的 .jpg和 .txt前面的名称要一样，比如apple_01.jpg 和 apple_01.txt ）

建立一个 dataset 文件夹，并在该目录下建立 images 和 labels 文件夹，并在这两个文件夹下面创建两个子文件夹，一个叫 train 用于存储训练集，一个叫 val 用于存储验证集

比如标签的训练集.txt文件就放在 label\val

接着我们就用labelimg框选训练集

🧭 一、启动软件

如果你是通过命令行运行的，输入：
```
python labelImg.py
```
或者直接运行 labelImg.exe（Windows 下）。

🖼 二、打开图片文件夹

点击左上角的按钮：

Open Dir：选择你存放图像的文件夹（比如 images/）
如果 Open Dir 找不到的话，也可以直接 Open 打开图片

🧩 三、选择保存标注的格式

点击：

YOLO（保存为 .txt 格式，适用于 YOLO 模型）

✏️ 四、开始标注步骤

点击工具栏上的 Create RectBox 按钮（或快捷键 w）
用鼠标在图片上 点击并拖动，画出一个框框（比如把猫框住）
系统会弹出一个窗口，输入标签名（如：cat） → 回车
你可以继续框其他对象，每个对象一个框

💾 五、保存标注

点击左上角的：

Save 按钮（或快捷键 Ctrl+S），系统会把标注保存为 .xml 或 .txt 文件，放在和图像同目录下。

⏭ 六、切换到下一张图像

点箭头按钮 ▶ 或按 d 切换到下一张图像继续标注
按 a 切换上一张图像
按 Ctrl + S 随时保存标注结果

⌨️ 常用快捷键：

w 框选

Ctrl + s 保存

d 下一张

a 上一张

Delete 删除框

Crtl + e 修改标签名

然后我们写 fruits.yaml # 水果识别数据集配置（注意是文件中写的是相对路径）

# 训练和验证数据集路径（相对于YOLOv5目录）
path: ../dataset  # 数据集根目录
train: images/train  # 训练图像路径（相对于path）
val: images/val  # 验证图像路径（相对于path）# 类别
nc: 5  # 类别数量
names: ['apple', 'banana', 'orange', 'strawberry', 'kiwi']  # 类别名称

迁移训练

运行训练脚本 train_fruis.py

import os
import argparse
import torch
import shutildef train_yolo(epochs=100, batch_size=16, img_size=640, weights='yolov5s.pt'):"""使用YOLOv5进行水果识别的迁移学习"""print(f"开始训练水果识别模型，训练轮次：{epochs}，批次大小：{batch_size}，图像尺寸：{img_size}")# 确保我们在正确的工作目录os.chdir(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'yolov5'))# 训练命令train_cmd = f"python train.py --img {img_size} --batch {batch_size} --epochs {epochs} --data ../dataset/fruits.yaml --weights {weights} --cache"# 执行训练print(f"执行命令: {train_cmd}")os.system(train_cmd)# 复制训练结果到models目录last_model = os.path.join('runs', 'train', 'exp', 'weights', 'last.pt')best_model = os.path.join('runs', 'train', 'exp', 'weights', 'best.pt')if os.path.exists(last_model):print("复制模型到models目录...")os.makedirs('../models', exist_ok=True)shutil.copy(last_model, '../models/fruits_last.pt')if os.path.exists(best_model):shutil.copy(best_model, '../models/fruits_best.pt')print("训练完成！模型已保存到models目录。")def parse_arguments():parser = argparse.ArgumentParser(description='YOLO水果识别迁移学习')parser.add_argument('--epochs', type=int, default=100, help='训练轮次')parser.add_argument('--batch-size', type=int, default=16, help='批次大小')parser.add_argument('--img-size', type=int, default=640, help='图像尺寸')parser.add_argument('--weights', type=str, default='yolov5s.pt', help='预训练权重 (yolov5s.pt, yolov5m.pt, yolov5l.pt, yolov5x.pt)')return parser.parse_args()if __name__ == "__main__":args = parse_arguments()train_yolo(epochs=args.epochs, batch_size=args.batch_size, img_size=args.img_size, weights=args.weights)

继续训练

如果新增了训练集的话，直接运行上面的 train_fruits.py 会生成新的模型

如果我们想继续训练我们之前的模型

import os
import argparse
import shutildef continue_training(epochs=50, batch_size=16, img_size=640, weights='models/fruits_best.pt'):"""使用已训练的水果模型继续训练（增量学习）"""print(f"使用已训练的模型继续训练水果识别模型，训练轮次：{epochs}，批次大小：{batch_size}")# 确保我们在正确的工作目录os.chdir(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'yolov5'))# 检查指定的权重文件是否存在weights_path = os.path.join('..', weights) if not os.path.isabs(weights) else weightsif not os.path.exists(weights_path):print(f"错误：找不到指定的权重文件：{weights_path}")print("如果这是您第一次训练，请使用train_fruits.py脚本，或者指定yolov5s.pt作为起点")return False# 训练命令 - 使用之前训练的模型继续训练train_cmd = f"python train.py --img {img_size} --batch {batch_size} --epochs {epochs} --data ../dataset/fruits.yaml --weights {weights_path} --cache"# 执行训练print(f"执行命令: {train_cmd}")os.system(train_cmd)# 复制训练结果到models目录last_model = os.path.join('runs', 'train', 'exp', 'weights', 'last.pt')best_model = os.path.join('runs', 'train', 'exp', 'weights', 'best.pt')if os.path.exists(last_model):print("复制更新后的模型到models目录...")os.makedirs('../models', exist_ok=True)shutil.copy(last_model, '../models/fruits_continued_last.pt')if os.path.exists(best_model):shutil.copy(best_model, '../models/fruits_continued_best.pt')print("增量训练完成！更新后的模型已保存到models目录。")return Truedef parse_arguments():parser = argparse.ArgumentParser(description='YOLO水果识别增量训练')parser.add_argument('--epochs', type=int, default=50, help='训练轮次')parser.add_argument('--batch-size', type=int, default=16, help='批次大小')parser.add_argument('--img-size', type=int, default=640, help='图像尺寸')parser.add_argument('--weights', type=str, default='models/fruits_best.pt', help='之前训练好的模型权重文件路径')return parser.parse_args()if __name__ == "__main__":args = parse_arguments()continue_training(epochs=args.epochs, batch_size=args.batch_size, img_size=args.img_size, weights=args.weights)