初めに

公式にて紹介されている物体検出(detection)のTransfer Learningを行う際、jetson-nano上で学習を行うと処理時間がかかるためNVIDIA Geforce GTX 1060が搭載されたPCで学習を行って結果のみをjetson-nanoに移動させて画像認識を行いました。その時の方法をメモします。

動作環境

jetson nano(jetpack 4.4をインストールした)
ubuntu20.04(NVIDIA Geforce GTX 1060が搭載されたPC)
- pytorch 1.4

NVIDIA Geforce GTX 1060が搭載されたPC上での準備

学習用環境を用意

参考と同様にpython3.6+pytorchの環境を用意する。
ただしpytorch 1.4をインストールする(1.6だとだめでした)

conda install pytorch==1.4.0 torchvision==0.5.0 cudatoolkit=10.1 -c pytorch

その他必要なものをインストール

pip install opencv-python
conda install pandas pillow matplotlib

学習用コードの準備

コンピュータのアーキテクチャが異なるため(?)jetson-inferenceをそのままbuildすることができません。しかし学習を行うだけならソースコードだけで十分なのでコードのみ拝借します

git clone --recursive https://github.com/dusty-nv/jetson-inference

モデルなど必要なもののダウンロード

公式通り

download

cd jetson-inference/python/training/detection/ssd
wget https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O models/mobilenet-v1-ssd-mp-0_675.pth
pip3 install -v -r requirements.txt

をおこなう。
※ boto3がないと言われたらconda install -c anaconda boto3でインストールする。参考

データの用意

vottでのアノテーションデータの作成を参考にアノテーションを行う。また、Pacal VOC形式でエクスポートする。
できた[任意の名前1]-PascalVOC-exportフォルダをjetson-inference/python/training/detection/ssd/に入れる
クラスごとにファイルが[任意の名前1]-PascalVOC-export/ImageSets/Mainにできているがその中から一つのみ残して各行末の1,-1,と.jpgを削除
上のファイルを適当に振り分けtest.txt,train.txt,trainval.txt,val.txtとして保存する。(train.val,testの3つに分けtrainvalはtrainとvalの内容を足したものにする)
参考 PASCAL VOC③（Datasetの確認）｜機械学習の有名データセットや評価指標を確認する #3
jetson-inference/python/training/detection/ssd/vision/datasets/voc_dataset.pyの46行目のクラスをつけたtagと同じにする。
参考 https://github.com/dusty-nv/jetson-inference/issues/674

学習の実行

train

python3 train_ssd.py --dataset-type=voc --batch-size=16 --data=[任意の名前1]-PascalVOC-export --model-dir=models/[任意の名前2]

で学習を行う(--batch-size=16を指定しないとRuntimeError: invalid argument 2: non-empty vector or matrix expected at…というエラーが出ました(理由は解決していません))

jetson-nanoでの準備

jetson-inferenceのインストール

公式通りにjetson-nano上にインストールしていきます。

jetson-inference

sudo apt-get update
sudo apt-get install git cmake libpython3-dev python3-numpy
git clone --recursive https://github.com/dusty-nv/jetson-inference
cd jetson-inference
mkdir build
cd build
cmake ../
make -j$(nproc)
sudo make install
sudo ldconfig

途中でインストールしたいモデルやpytorchのバージョンが聞かれますが、モデルは自分が使いたいモデルを選択しpytorchはpython3.6対応のほうをインストールするようにしてください。

モデルなど必要なもののダウンロード

公式通り

download

cd jetson-inference/python/training/detection/ssd
wget https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O models/mobilenet-v1-ssd-mp-0_675.pth
pip3 install -v -r requirements.txt

をおこなう。

onnxファイルの作成

NVIDIA Geforce GTX 1060が搭載されたPCで学習した結果できた[任意の名前2]フォルダをjetson-nanoのjetson-inference/python/training/detection/ssd/modelsにコピペする。
公式通り

python3 onnx_export.py --model-dir=models/[任意の名前2]

を行う。

webカメラからの動画に対して認識を行う

これもまた公式通り。

webカメラをjetson-nanoに接続する。
v4l2の場合、以下を実行する。

detectnet --model=models/[任意の名前2]/ssd-mobilenet.onnx --labels=models/[任意の名前2]/labels.txt \--input-blob=input_0 --output-cvg=scores --output-bbox=boxes /dev/video0

カメラについては公式(カメラについて)を参考にしてください。

参考