AI打遊戲-肆(模型訓練)

發表於2023-09-22

目標

  • pip安裝依賴
  • (可選)安裝GPU版pytorch
  • 自定義訓練配置,模型訓練

步驟

文件

pip安裝依賴

# 安裝依賴
pip install ultralytics

# 耗時有點長,可以使用國內映象安裝
pip install -i https://mirrors.aliyun.com/pypi/simple/ ultralytics 
  • 測試效果

    • 隨便搜尋下載一張貓或狗的圖片,把路徑作為source引數
  • "yolov8n.pt"是官方預訓練好的模型,第一次使用會自動從Github上下載

    • 國內訪問Github網路不穩定,如果失敗多重試幾次,或自己下載放到根目錄中
yolo predict model=yolov8n.pt source='dog.jpg'

3_0_dog.jpg

(可選)重灌GPU版pytorch

  • ultralytics預設安裝的是cpu版pytorch

    • 導致使用英偉達顯示卡並安裝了CUDA和cuDNN,但訓練時依然是CPU
  • 執行下面的python檔案,檢視列印內容就可以看出來
import torch

# 判斷cuda是否生效
print(torch.cuda.is_available())

# 檢視torch版本
print(torch.__version__)
  • pip解除安裝CPU版本pytorch
  • 開啟pytorch官網,重新安裝GPU版pytorch
# 解除安裝cpu版本的pytorch
pip uninstall torch torchvision torchaudio

# 重新安裝gpu版本(cuda11.7)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
  • 重新執行上面的python檔案,看到版本從"2.0.1+cpu"變為"2.0.1+cu117"

自定義訓練配置

  • 參考ultralytics專案的coco128.yaml訓練配置,編寫自定義訓練配置
  • 將訓練集根目錄設定為上一期影片,從標註平臺匯出的目錄
  • 並將train、val路徑設定為images(標註檔會自動識別為"labels")
  • 示例如下:game.yml
path: E:\ai-play-game\project-1-at-2023-09-13-17-05-6275bec0  # dataset root dir
train: images  # train images (relative to 'path') 128 images
val: images  # val images (relative to 'path') 128 images

# Classes,最好使用英文,opencv支援中文比較麻煩
names:
  0: hellhound  # 地獄犬
  1: samurais   # 武士
  2: player     # 玩家
  3: fireflies  # 螢火蟲

訓練

  • 詳細訓練引數看官網ultralytics - train

    • data: 指定訓練配置路徑
    • model: 使用"yolo8n.yaml"模型(定義模型網路,"yolov8n.yaml"是內建yolov8n網路結構)
    • batch: 批大小,數值太大電腦扛不住
    • epoches: 訓練輪次,遊戲元素比較簡單,訓練一兩百次就有不錯的效果了
yolo detect train data=game.yml model=yolov8n.yaml batch=4 epochs=200

訓練結果驗證

  • 訓練時,自動在根目錄建立runs/detect/train目錄,可以看到各種訓練資料
  • 訓練結束後,"runs/detect/train/weights/"目錄裡是模型結果

    • best.pt,是本輪訓練中效果最好的模型
    • last.pt,是最後一次訓練的模型
  • 使用任意遊戲截圖測試效果
yolo predict model=runs/detect/train/weights/best.pt  source='E:/project-e8d01b35/images/xxx.jpg'

3_1_image.jpg

相關文章