目標
- 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'
(可選)重灌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'