クラス分類のTransfer Learningをpcで学習しjetson nanoにのせる

初めに

公式にて紹介されているクラス分類の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(1.6だと動かないかもしれない)

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

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

学習用環境を用意

参考と同様にpython3.6+pytorchの環境を用意する。
ただし上に述べたようにpytorch 1.4をインストールする

学習用コードの準備

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

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

学習データの準備

1. 自前の画像データを集めjetson-inference/python/training/classification/にmy_dataフォルダを作成してその中にtrain/[クラス名]/[jpgの画像データ],val/[クラス名]/[jpgの画像データ],test/[クラス名]/[jpgの画像データ]とする。
2. 先程作ったmy_dataの中にlabals.txtを作成し学習するクラスを列挙する。(書式は公式にあります.)

学習の実行

train

cd jetson-inference/python/training/classification
python train.py --model-dir=[学習後のモデルの出力先] --arch=[使いたいモデル] my_data
#どのようなモデルが使えるのかなど詳細はpython train.py --helpを実行して確かめる

を実行する。

onnxファイルの作成

公式通り

python3 onnx_export.py --model-dir=[学習後のモデルの出力先]

を実行する。

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対応のほうをインストールするようにしてください。

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

これもまた公式通り。

1. webカメラをjetson-nanoに接続する。
2. NVIDIA Geforce GTX 1060が搭載されたPCで学習した結果できた[学習後のモデルの出力先]フォルダをjetson-nanoのjetson-inference/python/training/classification/にコピペする。
3. NVIDIA Geforce GTX 1060が搭載されたPCのmy_data/labels.txtをjetson-nanoの同じ場所にコピペする(必要なのはlabels.txtのみなので画像データはコピペしなくて良いです)
4. v4l2の場合

imagenet.py --model=[学習後のモデルの出力先]/[出力されたonnxファイルの名前].onnx --input_blob=input_0 --output_blob=output_0 --labels=my_data/labels.txt v4l2:///dev/video0

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

参考

nvidia公式によるjetson-inferenceのインストール解説
nvidia公式によるクラス分類の再学習の解説