Jetson(NX)でPyTorchとYOLOv5を動かす(Dockerを使用)

Dockerを使うのが早いです

PyTorch

NVIDIA L4T PyTorch | NVIDIA NGCを使います

sudo docker run -it --runtime nvidia --network host nvcr.io/nvidia/l4t-pytorch:r32.5.0-pth1.7-py3

sudo docker pull nvcr.io/nvidia/l4t-pytorch:r32.5.0-pth1.7-py3
sudo docker run -it --rm --runtime nvidia --network host nvcr.io/nvidia/l4t-pytorch:r32.5.0-pth1.7-py3

YOLOv5

ultralytics/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLiteからcloneします。

cd ~
git clone https://github.com/ultralytics/yolov5.git
cd yolov5

pipで必要なパッケージをいれます。

pip3 install -U pip #アップデートしないと入らないパッケージがありました。
sed -i '/torchvision/d' requirements.txt
#torchvision0.8.1が入らないので、`requirement.txt`から、torchvisionの行を消しました。
pip3 install -r requirements.txt

#サンプルプログラムを動かします
python3 detect.py --source ./data/images/bus.jpg --weights yolov5s.pt --conf 0.4

Dockerfile

FROM nvcr.io/nvidia/l4t-pytorch:r32.5.0-pth1.7-py3

WORKDIR /root

RUN git clone https://github.com/ultralytics/yolov5.git \
&& cd yolov5 \
&& sed -i '/torchvision/d' requirements.txt \
&& pip3 install -U pip \
&& pip3 install -r requirements.txt