JetsonNano でYOLO v5

Jetson NanoでYOLO v5を動かしてみました。

YOLO v5 って?

YOLO v5とは物体検出モデルのこと。YOLO系は他のモデル(Efficient detとか)と比較して軽量であり、Jetsonのようなシングルボードコンピュータでも動作できるのが強み。

個人的にYOLO v5がいいなと思ったのは以下の2点。

  • モデルサイズがS, M, L, Xと4つも用意してある
  • 学習がとても楽(水増しとかもコンフィグを変更するだけで使える)

今回はJetson Nano上で推論のみをおこなっていく。

github.com

前準備

今回DockerとCUDAが必要なので、Jetpackをいれておくこと。

  • L4T:R32.4.4 (まあ、Dockerで動かすので最新ので問題ない)
  • L4T:R32.5

Docker で動かすから問題ないと考えていたが、DockerイメージのL4TバージョンとホストのL4Tバージョンが一致していないと動かない。現在Pytorch1.7以上のイメージは、L4T R32.5以上でしか配布していないので、R32.5以上をインストールするか、バージョンアップしておくこと。

準備

YOLO v5はPytorchで動いているので、PytorchのDockerイメージを起動する。なければ 勝手にダウンロードしてくれる。Pytorchは1.7以上が必要なので1.7以上のイメージをインストールすること。

$ docker run --runtime nvidia -it --net host nvcr.io/nvidia/l4t-pytorch:r32.5.0-pth1.7-py3 bash

上のコマンドを実行するとDockerコンテナ内に移動するので、必要なパッケージ、ライブラリをインストールする。SSHはいらないが、推論結果をやりとりするのに便利なのでいれる。 YOLO v5 のrequire.txtを使用して入れると、OpenCVが動かないので下記コマンドで入れること。

$ apt update && apt install -y python3-opencv ssh && pip3 install -U pip && pip3 install requests tqdm pyyaml seaborn pycocotools thop

最後にYOLO v5 をクローンしてくる。

$ git clone https://github.com/ultralytics/yolov5.git

今回Jetson Nano用にパッケージを入れているので、バージョンは一致しない。YOLOv5ではバージョンチェックが入るのでコメントアウトしておく。

$ vim requirements.txt

以下をコメントアウト

#opencv-python>=4.1.2
#tensorboard>=2.2
#torchvision>=0.8.1

実行

サンプルとして用意されている画像を推論してみる。 まずはモデルサイズが最も小さい、yolov5sを試してみる。

/yolov5# python3 detect.py --source data/images --weights yolov5s.pt --conf 0.25
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, img_size=640, iou_thres=0.45, name='exp', nosave=False, project='runs/detect', save_conf=False, save_txt=False, source='data/images', update=False, view_img=False, weights=['yolov5s.pt'])
YOLOv5 \U0001f680 v4.0-170-g866bc7d torch 1.7.0 CUDA:0 (NVIDIA Tegra X1, 3964.1328125MB)

Fusing layers... 
Model Summary: 224 layers, 7266973 parameters, 0 gradients, 17.0 GFLOPS
image 1/2 /yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, Done. (0.247s)
image 2/2 /yolov5/data/images/zidane.jpg: 384x640 2 persons, 1 tie, Done. (0.158s)
Results saved to runs/detect/exp3
Done. (0.938s)

推論にかかった時間は0.247秒と0.158秒であった。リアルタイム推論も行けそうな感じではある。

他のサイズのモデルも試してみる。

  • M
/yolov5# python3 detect.py --source data/images --weights yolov5m.pt --conf 0.25
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, img_size=640, iou_thres=0.45, name='exp', nosave=False, project='runs/detect', save_conf=False, save_txt=False, source='data/images', update=False, view_img=False, weights=['yolov5m.pt'])
YOLOv5 \U0001f680 v4.0-170-g866bc7d torch 1.7.0 CUDA:0 (NVIDIA Tegra X1, 3964.1328125MB)

Downloading https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5m.pt to yolov5m.pt...
100%|##########################################################################################| 41.1M/41.1M [00:16<00:00, 2.57MB/s]

Fusing layers... 
Model Summary: 308 layers, 21356877 parameters, 0 gradients, 51.3 GFLOPS
image 1/2 /yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, Done. (0.605s)
image 2/2 /yolov5/data/images/zidane.jpg: 384x640 2 persons, 1 tie, Done. (0.772s)
Results saved to runs/detect/exp4
Done. (2.603s)

0.605秒と0.772秒

  • L
/yolov5# python3 detect.py --source data/images --weights yolov5l.pt --conf 0.25
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, img_size=640, iou_thres=0.45, name='exp', nosave=False, project='runs/detect', save_conf=False, save_txt=False, source='data/images', update=False, view_img=False, weights=['yolov5l.pt'])
YOLOv5 \U0001f680 v4.0-170-g866bc7d torch 1.7.0 CUDA:0 (NVIDIA Tegra X1, 3964.1328125MB)

Downloading https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5l.pt to yolov5l.pt...
100%|##########################################################################################| 90.2M/90.2M [00:26<00:00, 3.60MB/s]

Fusing layers... 
Model Summary: 392 layers, 47025981 parameters, 0 gradients, 115.4 GFLOPS
image 1/2 /yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, 1 potted plant, Done. (1.129s)
image 2/2 /yolov5/data/images/zidane.jpg: 384x640 3 persons, 2 ties, Done. (0.953s)
Results saved to runs/detect/exp5
Done. (3.817s)

1.129秒と0.953秒

  • X
/yolov5# python3 detect.py --source data/images --weights yolov5x.pt --conf 0.25
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, img_size=640, iou_thres=0.45, name='exp', nosave=False, project='runs/detect', save_conf=False, save_txt=False, source='data/images', update=False, view_img=False, weights=['yolov5x.pt'])
YOLOv5 \U0001f680 v4.0-170-g866bc7d torch 1.7.0 CUDA:0 (NVIDIA Tegra X1, 3964.1328125MB)

Downloading https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5x.pt to yolov5x.pt...
100%|############################################################################################| 168M/168M [00:56<00:00, 3.11MB/s]

Fusing layers... 
Model Summary: 476 layers, 87730285 parameters, 0 gradients, 218.8 GFLOPS
image 1/2 /yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, 1 potted plant, Done. (2.373s)
image 2/2 /yolov5/data/images/zidane.jpg: 384x640 2 persons, 3 ties, Done. (1.819s)
Results saved to runs/detect/exp6
Done. (6.118s)

2.373秒と1.819秒。 ひとまず、どのサイズのモデルであっても、動かすことは可能な模様。

実行(失敗バージョン)

なお、これはL4T R32.4で動かしたときの動作であり、GPUドライバーが動かないとのことでCPUでの推論になっている。参考までに残しておく。

/yolov5# python3 detect.py --source data/images --weights yolov5s.pt --conf 0.25
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, img_size=640, iou_thres=0.45, name='exp', project='runs/detect', save_conf=False, save_txt=False, source='data/images', update=False, view_img=False, weights=['yolov5s.pt'])
/usr/local/lib/python3.6/dist-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /media/nvidia/WD_NVME/PyTorch/JetPack_4.4.1/pytorch-v1.7.0/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
YOLOv5 v4.0-121-gba18528 torch 1.7.0 CPU

Fusing layers... 
Model Summary: 224 layers, 7266973 parameters, 0 gradients, 17.0 GFLOPS
image 1/2 /yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, Done. (1.417s)
image 2/2 /yolov5/data/images/zidane.jpg: 384x640 2 persons, 1 tie, Done. (1.104s)
Results saved to runs/detect/exp2
Done. (2.762s)

推論処理に1.4秒と1.1秒かかっている。

他のモデルも使ってみた。

  • M
/yolov5# python3 detect.py --source data/images --weights yolov5m.pt --conf 0.25
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, img_size=640, iou_thres=0.45, name='exp', project='runs/detect', save_conf=False, save_txt=False, source='data/images', update=False, view_img=False, weights=['yolov5m.pt'])
/usr/local/lib/python3.6/dist-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /media/nvidia/WD_NVME/PyTorch/JetPack_4.4.1/pytorch-v1.7.0/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
YOLOv5 v4.0-121-gba18528 torch 1.7.0 CPU

Downloading https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5m.pt to yolov5m.pt...
100%|##########################################################################################| 41.1M/41.1M [00:15<00:00, 2.87MB/s]

Fusing layers... 
Model Summary: 308 layers, 21356877 parameters, 0 gradients, 51.3 GFLOPS
image 1/2 /yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, Done. (2.956s)
image 2/2 /yolov5/data/images/zidane.jpg: 384x640 2 persons, 1 tie, Done. (2.315s)
Results saved to runs/detect/exp3
Done. (5.412s)
  • L
/yolov5# python3 detect.py --source data/images --weights yolov5l.pt --conf 0.25
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, img_size=640, iou_thres=0.45, name='exp', project='runs/detect', save_conf=False, save_txt=False, source='data/images', update=False, view_img=False, weights=['yolov5l.pt'])
/usr/local/lib/python3.6/dist-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /media/nvidia/WD_NVME/PyTorch/JetPack_4.4.1/pytorch-v1.7.0/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
YOLOv5 v4.0-121-gba18528 torch 1.7.0 CPU

Downloading https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5l.pt to yolov5l.pt...
100%|##########################################################################################| 90.2M/90.2M [00:38<00:00, 2.44MB/s]

Fusing layers... 
Model Summary: 392 layers, 47025981 parameters, 0 gradients, 115.4 GFLOPS
image 1/2 /yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, 1 potted plant, Done. (5.590s)
image 2/2 /yolov5/data/images/zidane.jpg: 384x640 3 persons, 2 ties, Done. (4.408s)
Results saved to runs/detect/exp4
Done. (10.148s)
  • X
/yolov5# python3 detect.py --source data/images --weights yolov5x.pt --conf 0.25
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, img_size=640, iou_thres=0.45, name='exp', project='runs/detect', save_conf=False, save_txt=False, source='data/images', update=False, view_img=False, weights=['yolov5x.pt'])
/usr/local/lib/python3.6/dist-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /media/nvidia/WD_NVME/PyTorch/JetPack_4.4.1/pytorch-v1.7.0/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
YOLOv5 v4.0-121-gba18528 torch 1.7.0 CPU

Downloading https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5x.pt to yolov5x.pt...
100%|############################################################################################| 168M/168M [01:00<00:00, 2.93MB/s]

Fusing layers... 
Model Summary: 476 layers, 87730285 parameters, 0 gradients, 218.8 GFLOPS
image 1/2 /yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, 1 potted plant, Done. (8.954s)
image 2/2 /yolov5/data/images/zidane.jpg: 384x640 2 persons, 3 ties, Done. (7.382s)
Results saved to runs/detect/exp5
Done. (16.482s)