3. 学習を行う(GPU搭載のPC上で行う場合)

次に、先程作成したアノテーション付き画像を学習させ、学習結果のモデルファイルを出力します。 こちらでは学習をLinux PC上で行う方法を説明します。 nvidia GPUが使えるLinux PCを所持している場合は、こちらの方法で学習させた方が学習時間が短く、安定しています。

推奨環境

  • nvidia GPUが搭載されたPC(Ubuntu 22.04で動作確認済み。cudaをインストールしておくこと。)

事前準備

以下すべてPC上での作業となります。
  1. PCのターミナル上で、akari_yolo_training_localのレポジトリをgitからcloneします。

git clone https://github.com/AkariGroup/akari_yolo_training_local.git
  1. 2. 画像のアノテーションをする で準備した画像とアノテーションファイルが含まれるimagesディレクトリを、先程Cloneした akari_yolo_training_local/images 内にコピーします。

    下図のディレクトリ構成になっていることを確認してください。
../../../_images/train_directory_03.jpg
  1. pythonがインストールされていない場合は、インストールします。

sudo apt install python3.10 python3.10-venv
  1. akari_yolo_training_localディレクトリに移動し、submoduleの更新、仮想環境の作成をします(初回のみ必要)。

cd akari_yolo_training_local
git submodule update --init
python -m venv venv
  1. 仮想環境を有効にします。

. venv/bin/activate
  1. 必要なライブラリをインストールします。

pip install -r requirements.txt
  1. データセットの画像、アノテーションファイルの8割を学習に、2割をバリデーション用に振り分けます。下記pythonを実行することで、自動で行われます。

python3 divide_dataset.py
  1. 学習用のパラメータファイルを作成します。下記pythonを実行することで、 akari_train_data.yaml が自動で作成されます。

python3 make_train_data.py
  1. 学習済みモデルの重みを取得します。

wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt -P data/
  1. 学習用のコンフィグファイルを作成します。下記pythonを実行することで、 akari_train_cfg.yaml が自動で作成されます。

python3 make_train_cfg.py

学習の実行

下記コマンドを実行すると、学習を開始します。 学習が終了するまで、しばらく待ちます。

python3 yolov7/train.py --epochs 300 --workers 8 --device 0 --batch-size 16 --data data/akari_train_data.yaml --img 640 640 --cfg data/akari_train_cfg.yaml --weights data/yolov7-tiny.pt --hyp yolov7/data/hyp.scratch.tiny.yaml --name result --exist-ok

注釈

もし undefined symbol: __nvJitLinkAddData_12_1, version libnvJitLink.so.12 というエラーが出た場合は、自身のvenv内のnvidiaライブラリにLD_LIBRARY_PATHを通す必要があります。

export LD_LIBRARY_PATH=~/"akari_yolo_training_localのパス"/venv/lib64/python3.10/site-packages/nvidia/nvjitlink/lib:$LD_LIBRARY_PATH

テストの実行

学習が終了すると、akari_yolo_training_localディレクトリ直下にrunsディレクトリが作成され、 runs/train/result 内に学習モデル(.pt)が保存されます。
best.pt が学習スコアが最も良かったファイルとなるので、こちらのファイルでテストを実施してみます。
下記コマンドを実行すると、学習したモデルを用いてテストを開始します。
python3 yolov7/test.py --data data/akari_train_data.yaml --iou-thres 0.65 --device 0 --weights runs/train/result/weights/best.pt --task test --name result --exist-ok
テスト結果は、 akari_yolo_training_local/runs/test/result に保存されます。
テスト画像に対する推論結果が画像として保存されているので、期待通り検出されているかを確認してください。

学習モデルの完成

テスト結果が問題なければ、 akari_yolo_training_local/runs/train/result/best.pt を次のモデル変換で使います。

3. モデル変換を行う へ進む

2. 画像のアノテーションをする へ戻る