Windows11,Anaconda3上でYOLOv5を動かしてみる(GPU編)

目的

画像データから物体認識ができる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へ置いた。

• 動画ファイルを読み込み物体認識を行った結果をフレーム毎にJPEGで保存する
  https://github.com/roxa-delphi/yolov5_test/blob/main/detectedFrameImage.py

• 動画ファイルを読み込み物体認識を行った結果を動画ファイルへ保存する
  https://github.com/roxa-delphi/yolov5_test/blob/main/detectedMovie.py

参考URL

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