ゆるエンジニアはいろいろ遊びたい

FAエンジニアが週末にいろいろ遊ぶブログです

YOLOをGPUで動かしてみた

ShelPiで使っている物体検出モデルのYOLOですが、現在CPUで動作しています。
これでも全然高速なのですが、今回はGPUで動かしてみたいと思います。
GPUで動作させるためには、pytorchというものを入れる必要があるようです。

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install ultralytics

インストールが完了したら、GPUを認識しているかチェックします。
pythonで

import torch
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))

結果

  • True
  • NVIDIA GeForce RTX3060(自分のGPU)

となれば成功
GPUで動かすには、modelにcudaを指定するといいようです。

model = YOLO("yolov8n.pt")
model.to("cuda")
#または
results = model(frame, divice="cuda")

GPUを使っているかの確認は

nvidia-smi

で、

====================================
|Voltage     Uncorr.ECC|
|GPU-Util Compute M.|
|                       MIG M.|
====================================
|                             N/A|
|       42%          Default|
|                             N/A|

こんな感じの画面が出てきますので、このGPU-Utilの部分がYOLO実行中に増加すればGPUを使っていると判断できます。
画像1枚ではすぐ終わってしまいわからないので、動画ファイルを与えてあげるとわかりやすいです。
YOLO自体がすごく軽いので、GPUの恩恵を全く感じませんでしたが、今後他のモデルを同時に使うときに役立ちそうです。