1 はじめに
CX 事業本部 delivery部の平内（SIN）です。
今年(2023年)5月に公開された YOLO-NASは、従来のYOLOv5,v8などより性能向上されたものとなっており、11月には、骨格検出モデルも追加され、オブジェクト検出モデルとして、今最も注目されているのではないでしょうか。
https://github.com/Deci-AI/super-gradients
ということで、Jetson AGX Orinで動作させて見たいと思ったのですが、ちょっと手こずったので、作業した記録として、この記事を投稿させて頂きました。
2 ImportError: cannot import name 'get_rank' from 'torch.distributed'
手こずったのは、ズバリ下記のエラーです。
# python3 -c "import super_gradients"
File "/home/super-gradients/src/super_gradients/training/utils/distributed_training_utils.py", line 12, in <module>
from torch.distributed import get_rank, all_gather_object
ImportError: cannot import name 'get_rank' from 'torch.distributed'Jetson AGX Orin では、公式で準備されているモジュールでは、CUDAが利用できません。
$ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
$ python3 -c 'import torch;print(torch.cuda.is_available())'
Falseこのため、通常、Nvidiaで準備されているWhlや、Dockerイメージを利用するのですが、こちらのPyTorchモジュールでは、先のエラーとなって、YOLO-NASが利用できませんでした。 https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch
エラーが出る原因は、Nvidiaで配布しているPyTorchモジュールが、USE_DISTRIBUTED=0でビルドされているためだと思います。古い配布をたどると、DISTRIBUTEDが有効になったものも有るようなのですが、ちょっとPyTorchのバージョンが古くなってしまうため、今回は、自前でビルドすることにしました。
3 PyTorchのビルド
(1) nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3
Nvidiaで配布されているDockerイメージを、作業の起点としました。
$ xhost +
$ docker run -it --rm --runtime nvidia --shm-size=1g -v $(pwd)/home:/home -e DISPLAY=:0 --network host nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3
#(2) build
最初にビルドに必要なpyyamlをインストールしています。
# pip install pyyaml続いて、githubからダウンロードしたソースから PyTorch をビルドしますが、この時、USE_DISTRIBUTED=1 を設定しています。
# git clone https://github.com/pytorch/pytorch.git
# cd pytorch
# export USE_DISTRIBUTED=1
# export TORCH_CUDA_ARCH_LIST="8.7"
# export PYTORCH_BUILD_VERSION=2.1.1
# export PYTORCH_BUILD_NUMBER=1
# python3 setup.py bdist_wheel
[xxxx/65xx] 小一時間かかる・・・(3) install
ビルドが完了したら、distに作成されたwhlをインストールします
# cd dist
# pip install torch-2.1.1-cp38-cp38-linux_aarch64.whlバージョンを確認しています。cudaや、torch.distributedが有効になっていることも確認できます。
# python3 -c "import torch;print(torch.__version__)"
2.1.1
# python3 -c 'import torch;print(torch.cuda.is_available())'
True
# python3 -c "import torch;print(torch.distributed.is_available())"
True
# python3 -c "import torch; print( torch.cuda.get_device_name(0) )"
'Orin'(4) torchvision
YOLO-NASでは、torchvisionも必要です。PyTorch 2.1 に対応するtorchvisionは 0.16なので、これをインストールしました。
# git clone --branch release/0.16 https://github.com/pytorch/vision torchvision
# cd torchvision
export BUILD_VERSION=0.16.0
# python3 setup.py install --user
# python3 -c "import torchvision; print( torchvision.__version__ )"
0.16.0(5) test
簡単に、動作確認しています。
# python3
Python 3.8.10 (default, Nov 14 2022, 12:59:47)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> device = torch.device('cuda', 0)
>>> torch.rand(4).to(device)
tensor([0.2048, 0.9094, 0.4167, 0.9344], device='cuda:0')4 super-gradientsのインストール
pip install super-gradientsでは、うまくインストールできなかったの、Githubからダウンロードしてインストールしました。
最新は、2023年11月23日リリースの3.5.0です。 https://pypi.org/project/super-gradients/#history
(1) cv2の無効化
Nvidiaで配布されたDockerイメージでは、最初からOpenCVがセットアップされているのですが、pipで管理されていないため、pip installで競合してしまいます。そこで、先にcv2を無効化しておきます。
# mv /usr/lib/python3.8/dist-packages/cv2 /usr/lib/python3.8/dist-packages/cv2.bak
# mv /usr/local/lib/python3.8/dist-packages/cv2 /usr/local/lib/python3.8/dist-packages/cv2.bak
# python3 -c 'import cv2'
ModuleNotFoundError: No module named 'cv2'(2) setuptools -U
setuptoolsのバージョン古くてエラーが出てしまうので、最新に更新します。
# pip install setuptools -U
# pip list| grep setuptools
setuptools 69.0.2(3) install
super-gradientsをGithubからダウンロードしてインストールします。
# git clone https://github.com/Deci-AI/super-gradients.git
# cd super-gradients
# pip install -r requirements.txt
# python3 setup.py install
# pip3 install -e .pytorch_quantizationｍのimportでWARNINGとなっていますが、一応、利用可能になりました。
# python3 -c "import super_gradients"
The console stream is logged into /root/sg_logs/console.log
INFO - crash_tips_setup.py - Crash tips is enabled. You can set your environment variable to CRASH_HANDLER=FALSE to disable it
WARNING - __init__.py - Failed to import pytorch_quantization
INFO - font_manager.py - generated new fontManager
WARNING - calibrator.py - Failed to import pytorch_quantization
WARNING - export.py - Failed to import pytorch_quantization
WARNING - selective_quantization_utils.py - Failed to import pytorch_quantization(4) test
画像表示のためpython3-tkをインストールします。
# apt-get install python3-tk物体検出モデルを試してみました。
# python3
>>> import super_gradients
>>> yolo_nas = super_gradients.training.models.get("yolo_nas_l", pretrained_weights="coco").cuda()
>>> yolo_nas.predict("office.jpg", conf=0.5).show()
5 Pase(骨格検出)
セットアップが完了したので、Poseを試してみています。 yolo_nas_pose_lで、1280×720の画像を処理する時間は、0.35〜0.45sec程度でした。
import os
import time
import super_gradients
yolo_nas = super_gradients.training.models.get(
"yolo_nas_pose_l", pretrained_weights="coco_pose"
).cuda()
input_path = "./input"
output_path = "./output"
if not os.path.exists(output_path):
os.makedirs(output_path)
files = os.listdir(input_path)
for i, file in enumerate(files):
if file.endswith(".jpg"):
time_sta = time.perf_counter()
model_predictions = yolo_nas.predict(
"{}/{}".format(input_path, file),
conf=0.5,
).save(output_path)
time_end = time.perf_counter()
tim = time_end - time_sta
print("{} {} {:.2f} sec".format(i, file, tim))
from_name = "{}/pred_0.jpg".format(output_path)
to_name = "{}/{}".format(output_path, file)
os.rename(from_name, to_name)
動画も試してみました。
from super_gradients.training import models
from super_gradients.common.object_names import Models
input_path = "./input_mp4"
output_path = "./output_mp4"
if not os.path.exists(output_path):
os.makedirs(output_path)
yolo_nas_pose = models.get("yolo_nas_pose_l", pretrained_weights="coco_pose").cuda()
files = os.listdir(input_path)
for i, file in enumerate(files):
if file.endswith(".mp4"):
yolo_nas_pose.predict("{}/{}".format(input_path, file), conf=0.5).save(
"{}/{}".format(output_path, file)
)6 最後に
今回は、Jetson AGX OrinでYOLO-NASを動作させて見ました。PyTorchのビルドが必要になりましたが、それ以外は、特に問題なく利用できているようです。
非常に高速に動作するYOLO-NASで、また、新たな楽しみが増えたかも知れません。
動画は、https://www.pexels.com/ja-jp/のものを、また、写真は、https://www.freebie-ac.jp/のものを利用させて頂きました。