PyTorch
SemanticSegmentation
ESPNet

ESPNetで自作データセットを学習してセグメンテーション

初めに

この記事はまだまだ未完であることをご容赦ください。

ESPNetを使うことになった

深層学習を使ったsemantic segmentationの研究は多々あるんですが
ロボットなどに実装するためリアルタイム性を求めた軽量なモデルはそう多くないです。
有名なので言えばENet・ICNet・ERFNetなど。
その他FCN・SegNet・Deeplabなどはかなり重量級なモデル。

現在のところリアルタイム性と分類精度はトレードオフですが
今回使ったESPNetはそこそこの精度でかなり軽量なモデルになっています。
論文によればJetson TX2上で9fps(1024✕512)という数字を叩き出しています。

今までCaffe上でENetを使用していたが精度に限界を感じたため
上位互換のESPNetを使ってみることにした。

ESPNetはPytorchで実装されており、初めてPytorchを使うにあたって
備忘録として記事を書きました。CaffeもCaffe2に移行しつつあることや
使い方に癖があるためこれをいい機会にPyTorchに乗り換えようか。

github:ESPNet
論文:ESPNet: Efficient Spatial Pyramid of Dilated Convolutions for Semantic Segmentation

動作環境

  • Ubuntu 16.04
  • Intel(R) Core(TM) i7-6800K CPU @ 3.40GHz
  • GeForce GTX 1080
  • CUDA 8.0
  • cuDNN 7.0
  • Anaconda 3-5.3.1(python3.7)
  • PyTorch 0.3.0
  • OpenCV 3.3.0

環境構築

pyenv・Anaconda・Pytorchの導入は以下の記事と同じですが一応書いておきます
Ubuntu 16.10にPyTorchを導入してみた

PyTorch初心者なので記事に従っていますが、PyTorchを入れる段階で
Anaconda入れたくない人はやり方があるので他の記事でおねがいします。

1. pyenv

terminal
git clone https://github.com/yyuu/pyenv.git ~/.pyenv

PATHを通すため.profileもしくは.bashrcなどに以下を追加

.bashrc
#pyenv
export PYENV_ROOT=$HOME/.pyenv
export PATH=$PYENV_ROOT/bin:$PATH
eval "$(pyenv init -)"

設定を反映させましょう

terminal
source ~/.bashrc
もしくは
source ~/.profile

2. Anaconda

Anacondaの最新バージョンを確認してインストール

terminal
pyenv install -l | grep anaconda3
pyenv install anaconda3-5.0.1

PATHを通すため.profileもしくは.bashrcに以下を追加

.bashrc
#Anaconda
export PATH="$PYENV_ROOT/versions/anaconda3-5.0.1/bin:$PATH"

設定を反映させましょう

terminal
source ~/.bashrc
もしくは
source ~/.profile

3. PyTorch

公式HP (http://pytorch.org) のTopで自身の環境を指定すると、
それに応じたインストールコマンドが指定されます。私の場合は

terminal
conda install pytorch=0.3.0 cuda80 -c pytorch

間違えて異なるバージョンをインストールしてしまった場合は

terminal
conda uninstall pytorch torchvision cuda80 -c soumith

ついでに必要なパッケージもインストール

terminal
conda install -c conda-forge opencv python-graphviz

実際にESPNetを使ってみる

自分の好きなところにESPNetのリポジトリをクローン

terminal
git clone https://github.com/sacmehta/ESPNet.git

ディレクトリはこんな構造になってます

ESPNet/
 ├ pretrained/    cityscapseで学習済みの重み
 │     ├ encoder/
 │     └ decoder/
 ├ sample_video/
 ├ test/        test用のプログラム
 │     ├ data/    テスト用のサンプル画像
 │     └ results/   VisualizeResult.pyを実行すると作成され出力結果がここに入る
 └ train/
       └ city/      cityscapesで学習する際に使うテキストファイル
        ├ train.txt トレーニングに用いるRGB画像とラベル画像のパスが羅列されたテキスト
        └ val.txt   バリーションに用いるRGB画像とラベル画像のパスが羅列されたテキスト

0. ちょっとテストしてみる

terminal
$ cd ESPNet/test
~/ESPNet/test$ python VisualizeResults.py --modelType 1 --p 2 --q 3

出力結果がtest/results以下に保存されていればオッケー

1. CityScapesで学習させる場合

1.1 train.txt/val.txtのパスを変更しなくて良い方法

Cityscapseのデータセットを下記のようにcity/以下に置く場合は1.2のパスを適応させる手順は必要ないです

ESPNet/
 ├ pretrained/
 ├ sample_video/
 ├ test/        
 └ train/
       └ city/
          ├ leftImg8bit/
          │       └...
          ├ gtFine/
          │       └...        
          ├ train.txt 
          └ val.txt   

1.2 train.txt/val.txtのパスを変更する必要がある場合

容量不足等の理由で外部ストレージにデータセットを置く場合は
train.txtとval.txtを適切なパスを示すように変更する必要があります。
テキストエディタとかで置換してやってください。

変更前
/leftImg8bit/train/zurich/zurich_000022_000019_leftImg8bit.png,/gtFine/train/zurich/zurich_000022_000019_gtFine_labelTrainIds.png
変更後
/home/…/dataset/cityscapes/image/train/zurich/zurich_000022_000019_leftImg8bit.png,/home/…/dataset/cityscapes/label/gtFine/train/zurich/zurich_000022_000019_gtFine_labelIds.png

また、ラベル画像のファイル名が元々のtrain.txt/val.txtでは
〜gtFine_labelTrainIds.png だったのですが実際の画像ファイル名は
〜gtFine_labelIds.png    で「Train」があるかないかの微妙な違いがありました
これも自分の環境に合わせて置換してください

私の用にルートからのフルパスで置換した場合は

ESPNet/train/loadData.pyの58,59行目
img_file = ((self.data_dir).strip() + '/' + line_arr[0].strip()).strip()
label_file = ((self.data_dir).strip() + '/' + line_arr[1].strip()).strip()
から
img_file = (line_arr[0].strip()).strip()
label_file = (line_arr[1].strip()).strip()

に変更して保存してください。
そうでない場合は学習実行時に引数として

python main.py --data_dir /home/media/ssd

の用にパスを指定してうまく行くようにしてください
しかしcityscapes自体は30classであるはずが
label画像を見ると30以上の値が入っているため、これをどうにかしなければいけない
私が作成したプログラムがあるのでパスやlabelのマップを変更してもらえれば使えると思います

2. 自分のデータセットで学習させる場合

2.1 RGB画像とラベル画像の用意

とりあえずどこでも良いのでファイル名が対応したRGB画像とラベル画像をこんな感じで用意してください

dataset/
 ├ train/
 │   ├ rgb/
 │   │   ├ 000000.png
 │   │   ├ 000001.png
 │   │   
 │   └ label/
 │       ├ 000000.png
 │       ├ 000001.png
 └ val/
     ├ rgb/
     │   ├ 000005.png
     │   ├ 000009.png
     │   
     └ label/
         ├ 000005.png
         ├ 000009.png

2.2 train.txtとval.txtの作成

ESPNet/train/にcreate_dataset_txt.pyを保存してください。
実行する前にtrain_rgb_dirなどは各自のデータセットを置いた場所に書き換えてください
train_txtとval_txtはESPNet/train/jisakuのようにESPNetがある場所したほうがわかりやすいと思います。

create_dataset_txt.py
# -*- coding: utf-8 -*-
import os
import glob

def main():
    Create_txt(train_rgb_dir, train_label_dir, train_txt)
    Create_txt(val_rgb_dir, val_label_dir, val_txt)

def Create_txt(rgb_dir, label_dir, txt_path):
    filelist_rgb = glob.glob(os.path.join(rgb_dir, '*.png'))
    txt = os.path.basename(txt_path)
    try:
        with open(txt_path, mode='x') as f:
            for file_path_rgb in filelist_rgb:
                filename = os.path.basename(file_path_rgb)
                file_path_label = os.path.join(label_dir, filename)
                s = file_path_rgb + ',' + file_path_label + '\n'
                f.write(s)

        print ('create :', txt)

    except FileExistsError:
        print(txt, 'already exists')

if __name__ == '__main__':
    #書き換えが必要です
    train_rgb_dir = '/path/to/dataset/train/rgb'
    train_label_dir = '/path/to/dataset/train/label'
    val_rgb_dir = '/path/to/dataset/val/rgb'
    val_label_dir = '/path/to/dataset/val/label'

    #jisakuの部分は好きな名前で構いません
    #※プログラムを実行する前にディレクトリを作成しておいてください
    train_txt = '/path/to/ESPNet/train/jisaku/train.txt'
    val_txt = '/path/to/ESPNettrain/jisaku/val.txt'

    main()

実行します

terminal
$ cd /ESPNet/train
~/ESPNet/train$ python create_dataset_txt.py

恐らく、jisakuディレクトリにtrain.txtとval.txtができていると思うので確認してください

ESPNet/
 ├ pretrained/
 ├ sample_video/
 ├ test/        
 └ train/
       ├ city/
       └ jisaku/
          ├ train.txt 
          └ val.txt   

2.3 ラベル画像のリマッピング※必要な人だけ

データセットに10クラスしか含まれていない場合でも
ラベル画像には20や30といったクラス数より大きい値が振られていることがある。
これに気づかずにクラス数を10に指定して学習させようとするよ
クラス数より大きい値がラベル画像中に出てきてしましまうため、エラーが起こる
これを解消するためにラベル画像を0から9(10クラスの場合)の値に収めるようにリマッピングする処理をする必要がある。

この先、追記予定

PS
別にリマッピングしなくてもラベル画像中に出てくる最大値をクラス数に指定すれば問題ないですが
私はなんか気持ち悪いのでラベル画像をリマッピングしています。

3. 学習スタート

学習は2段階あります。
①エンコーダーの学習
②デコーダーの学習

3.1 エンコーダーの学習

terminal
$ cd ESPNet/train
#Cityscapesでやる場合
ESPNet/train$ python main.py --scaleIn 8 --p 2 --q 8
#自作データセットの場合(--classesにはデータセットのクラス数を指定)
ESPNet/train$ python main.py --data_dir ./jisaku --classes 12 --scaleIn 8 --p 2 --q 8

・
・
・
[247/387] loss: 0.675 time:0.48
[248/387] loss: 0.651 time:0.48
[249/387] loss: 0.701 time:0.48
[250/387] loss: 0.536 time:0.48
[251/387] loss: 1.010 time:0.50
[252/387] loss: 0.862 time:0.48
[253/387] loss: 0.860 time:0.48
[254/387] loss: 0.758 time:0.48
[255/387] loss: 0.860 time:0.48
[256/387] loss: 0.794 time:0.48
[257/387] loss: 0.671 time:0.48
[258/387] loss: 0.738 time:0.48
・
・

重みのスナップショットは
ESPNet/train/results_enc__enc_2_8/
に保存されています。※2と8は学習時に指定したパラメータの値
収束してきたら、学習を止めて次はデコーダーの学習に移りましょう。

3.2 デコーダーの学習

エンコーダーの学習には引数は指定していませんでしたが、デフォルトでbatch_size=12を用いていました。
デコーダーではbatch_size=6なので指定してやりましょう。
環境によってはもっと小さい値をしてする必要もあるので、out of memoryと表示された場合は
コマンドラインで--batch_sizeを適切な値にしてやってください。

terminal
#Cityscapesでやる場合
ESPNet/train$ python main.py --batch_size 6 --scaleIn 1 --p 2 --q 8 --decoder True --pretrained ../pretrained/encoder/espnet_p_2_q_8.pth
#自作データセットの場合(--classesにはデータセットのクラス数を指定)
ESPNet/train$ python main.py --classes 12 --data_dir ./izunuma --batch_size 6 --scaleIn 1 --p 2 --q 8 --decoder True --pretrained ./results_enc__enc_2_8/model_81.pth

Encoder loaded!
Total network parameters: 354517
Data statistics
[131.36485 144.77011 134.60292] [76.74333  68.792656 71.73109 ]
[ 9.702483   9.986656   7.9102645  8.283521   3.7311087 10.492059
 10.206483   4.453087  10.42091   10.340107  10.338798   1.9798005]
Learning rate: 0.0005
[0/773] loss: 2.491 time:2.24
[1/773] loss: 2.485 time:0.40
[2/773] loss: 2.478 time:0.40
[3/773] loss: 2.478 time:0.40
[4/773] loss: 2.467 time:0.40
[5/773] loss: 2.465 time:0.40
・
・
・

学習終わり次第結果を追記予定

参考サイト

https://github.com/sacmehta/ESPNet
Ubuntu 16.10にPyTorchを導入してみた
Tensorflow Lite (TensorflowLite) / Tensorflow+RaspberryPi+Python で超軽量 "Semantic Segmentation" モデル "UNet" "ENet" を実装する軽量モデルその2