📚

Google ColabでYOLOv8を使ったモデル学習

に公開

環境

  • プラットフォーム: Google Colab(無料プラン)
  • GPU: T4(Colabの「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」から選択)
  • Pythonバージョン: 3.10以上(Colab標準)

追加のセットアップは不要で、Colab上でセルを順番に実行するだけでOKです。

セットアップ

ノートブックの作成

Google Colabを開いて新しいノートブックを作成。名前はなんでもいです。

GPUの設定

image.png

画面右からランタイプの変更を押して

image.png
ハードウェアアクセラレータからT4GPUを選択してください。
学習にGPUを用いるよ、という設定です。

以下のコードを1セルずつ実行します。

Ultralytics YOLOv8 をインストール

# YOLOv8をインストール
!pip install ultralytics

Jupyter NotebookやGoogle Colabでpipを実行するときは!をコマンドの前に付けます。
「これはPythonのコードではなく、シェルコマンドとして実行してください」と伝える役割があります。

インストールが完了したら

import ultralytics
print(ultralytics.__version__)

を実行してください。
8.3.174のようにバージョンが表示されたら成功しています。

YOLOv8が正常に動作しているか確認

from ultralytics import YOLO
model = YOLO("yolov8n.pt")  # Nanoモデルをロード
results = model("https://ultralytics.com/images/bus.jpg")  # サンプル画像で推論

# バウンディングボックス付きで表示
results[0].show()

上のプログラムを実行します。
与えるサンプル画像は以下のものです。
bus (1).jpg

このように、物体を検出できていれば成功です。
image.png

YOLOv8の学習

Aパターン:COCO128 で学習する(とりあえず試したい人向け)

サンプルデータセットのダウンロード

# Ultralytics の COCO128(128枚のCOCO小規模データ)
from ultralytics import YOLO
model = YOLO("yolov8n.pt")  # モデル(nano)

# COCO128はUltralytics側で自動ダウンロードされるので、手動で用意する必要はありません

学習の実行

model.train(
    data="coco128.yaml",  # 学習データ設定ファイル(YOLOに組み込み済み)
    epochs=10,            # エポック数(お試しなら10)
    imgsz=640,            # 画像サイズ
    batch=16              # バッチサイズ(GPUが小さい場合は8などにしてもOK)
)

簡単に解説すると

  • epochs:学習データ全体を何回繰り返して学習するかを指定(増やすと精度が上がる)
  • batch:一度にGPUに渡す画像の枚数(メモリに応じて調整)

学習結果はruns/detect/train/に出力されます。

image.png

学習後のモデル確認・ONNXエクスポート

学習が完了したら、best.ptというファイルが生成されます。これが学習済みのモデルです。

このbest.ptを、Jetsonなどのデバイスで使いやすいONNX形式に変換します。

# best.pt をONNX形式にエクスポート
model = YOLO('runs/detect/train/weights/best.pt')
model.export(format='onnx')

エクスポート結果:best.onnx(Jetson用にこのファイルを使う)

Bパターン:自分のデータセットで学習する場合

データセットの用意

https://qiita.com/Y_R_/items/07ced51178b38f7b2bde

この記事を参考に自作データセットを後輩に用意させた用意しました。
YOLO v8 PyTorchを学習に使用するライブラリとして選択しました。

詳しいことは別記事に書くと思います。

用意したデータセットのzipファイルをcolabにアップロードしたら

!unzip /content/用意したデータセット.zip -d /content/datasets/

でdatasetsフォルダ直下に解凍。

image.png

box_surfaceフォルダが自作のデータセットです。

/content/datasets/box_surface/
  ├── images/
  ├── labels/
  └── data.yaml

このようなフォルダ構成が表示されていればOKです。

Colabでの学習実行

model = YOLO('yolov8n.pt')  # 軽量モデル
model.train(data='/content/datasets/box_surface/data.yaml', epochs=30, imgsz=640)

完了した学習済みモデル(best.pt)は/content/runs/detect/train/weights/best.ptに保存されています。
今度はそれをONNX形式に変換していきます。

学習済みモデルをONNX形式でエクスポート

以下のコードでONNXファイルを出力します。

from ultralytics import YOLO

# best.pt(学習済みモデル)を読み込む
model = YOLO('runs/detect/train/weights/best.pt')

# ONNX形式でエクスポート
model.export(format='onnx')

終わりに

これで、Google Colabを使ったYOLOv8のモデル学習が完了しました。
学習したモデルを実際に使ってみてどうこう、という話はやったら記事にします。

Discussion