Google ColabでYOLOv8を使ったモデル学習
環境
- プラットフォーム: Google Colab(無料プラン)
- GPU: T4(Colabの「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」から選択)
- Pythonバージョン: 3.10以上(Colab標準)
追加のセットアップは不要で、Colab上でセルを順番に実行するだけでOKです。
セットアップ
ノートブックの作成
Google Colabを開いて新しいノートブックを作成。名前はなんでもいです。
GPUの設定
画面右からランタイプの変更
を押して
ハードウェアアクセラレータ
から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()
上のプログラムを実行します。
与えるサンプル画像は以下のものです。
このように、物体を検出できていれば成功です。
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/
に出力されます。
学習後のモデル確認・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パターン:自分のデータセットで学習する場合
データセットの用意
この記事を参考に自作データセットを後輩に用意させた用意しました。
YOLO v8 PyTorchを学習に使用するライブラリとして選択しました。
詳しいことは別記事に書くと思います。
用意したデータセットのzipファイルをcolabにアップロードしたら
!unzip /content/用意したデータセット.zip -d /content/datasets/
でdatasetsフォルダ直下に解凍。
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