目的

画像データから物体認識ができるYOLOv5をWindows上で試してみたので方法をまとめた。
前回はCPUのみで実行していたので、今回はGPUを使う方法を試してみた。

環境

CPU : Intel i7-10700F 8Core 16thread
GPU : NVIDIA GeForce RTX 3060 (Mem 12GB)
OS : Windows11
Anaconda : 3
Python : 3.9.7

YOLOv5のダウンロード

ここはCPU編と同じ。
https://github.com/ultralytics/yolov5 から git clone するか、zipファイルをダウンロードして展開しておく。

PyTorchのバージョンを確認する

https://pytorch.org/ のページへ行き、
Install ボタンをクリックすると
START LOCALY の画面が表示されるので、

PyTorch Build：Stable (1.10.1)
Your OS：Windows
Package：Conda
Language：Python
Computer Platform：CUDA 11.3

と選択すると下記が表示されるのでメモっておく。

Run this command : conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

Anacondaで環境作成

ANACONDA.NAVIGATORで仮想環境を新たに作成し、Terminalを開く。
そこでメモしておいたコマンドを下記のように実行することでPyTorchがインストールされる。

  conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

他にも必要なパッケージをインストールしておく

  pip install opencv-python
  pip install pandas
  pip install requests
  pip install PyYAML
  pip install tqdm
  pip install matplotlib
  pip install seaborn

実行してみる

CPU編の時に作成した test.py のディレクトリへ移動し、実行する。
ソースはこちら https://github.com/roxa-delphi/yolov5_test/blob/main/test.py

cd (target-dirctory)
python test.py

すると下記のようにメッセージが出て、認識したものにマーカーが付いた状態で動画が表示された。ちゃんとGPUが認識出来ていることが分かる。

YOLOv5  2022-1-22 torch 1.10.1 CUDA:0 (NVIDIA GeForce RTX 3060, 12288MiB)

60秒の動画の解析時間は

CPUのみ : 157秒
GPUあり : 40秒

と約3倍速いことを確認した。
GPUを使うことでリアルタイムに解析可能なことがわかった。

おまけ

下記の2つのスクリプトも作ったのでgithubへ置いた。

参考URL

Windows11,Anaconda3上でYOLOv5を動かしてみる(CPU編)
https://qiita.com/roxa/items/e735e8ceaeb326943674