Yolo遷移訓練-帶訓練源碼

下載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)