やること
特にひねりはなくタイトルの通りですが、OpenMMLab のページの ReadMe に書かれている通りにやるとできないので、できない部分を補足しています。SegFormer 以外を試す場合でもインストール周りはたぶん同じです。
2024/12現在、最低でも半年以上は ReadMe が修正されないままになってます
前提
- Ubuntu 22.04
- Cuda 12.2
- NVIDIA Driver Version 535.183.01 (Cuda のバージョンに合わせる）
cuda 12.5 などにする場合は、driver のバージョンをそれに合わせる必要があります。
cuda と driver のバージョンの合わせ方については、NVIDIA のダウンロードのページにいってインストーラを選択したときに、コマンドラインでダウンロードする時の deb ファイルの名前を見ます。
$ wget https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda-repo-ubuntu2204-12-2-local_12.2.0-535.54.03-1_amd64.deb
この行に 535.54 とあるので、cuda 12.2 のときは driver のバージョンを 535 にすれば大丈夫ということです。これ、単純に pytorch 使うだけならここまで厳密にあわせなくてもいいんですけど、これをあわせないと後述する mmcv のビルドが通らなくてはまったりします。
なお、以下では Anaconda を使用していますが、pip でも同様です。Ubuntu 24.04 は cuda と driver の整合性の問題で、mmcv のビルドが通らない可能性あり。
インストール
Pytorch のインストール
$ conda create -n mmseg$
$ conda activate mmseg
$ conda install python=3.8 -y
$ conda install pytorch torchvision -c pytorch
python は 3.8 の使用が絶対です（多分）。
Pytorch で cuda が使用可能かテスト
~$ python
Python 3.8.20 (default, Oct 3 2024, 15:24:27)
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True
>>>
True ならOKです。
MMSegmentation のインストール
$ pip install fsspec ftfy regex openmim
$ mim install mmengine
$ mim install "mmcv<2.2.0"
mmcv は公式の ReadMe だと mim install "mmcv>=2.2.0" のように書かれていますが、2.2.0 以上をインストールすると互換性がないと言われてエラーになります。mmcv を最初にインストールするときは、ビルド処理が入るのでかなり（マシンによって10～30分ほど）時間がかかります。2024/12 現在、上記コマンドでインストールすると mmcv はバージョン 2.1.0 が入る模様。
Cuda が 12.1 以下だったり、Cuda と NVIDIA driver のバージョンの整合性があっていない場合（冒頭に書いた通り）、nvidia-smi のバージョンと driver のバージョンが合わない場合などにもビルドエラーになります。
$ git clone https://github.com/open-mmlab/mmsegmentation.git
$ cd mmsegmentation
$ pip install -v -e .
以上でインストールは終了
学習テスト
データの用意
ADEChallengeData2016.zip を使う場合、データは HuggingFace や Kaggle から入手可能。
展開すると下記のようになります。ここでは mmsegmentation フォルダ直下に配置するものとします。
ADEChallengeData2016/
├── annotations
│ ├── training
│ └── validation
└── images
├── training
└── validation
dataset のコンフィグの用意
configs/_base_/datasets/ade20k.py を configs/_base_/datasets/ade20k_local.py にコピーする。
$ cp configs/_base_/datasets/ade20k.py configs/_base_/datasets/ade20k_local.py
configs/_base_/datasets/ade20k_local.py の data_root をデータを配置したパスに書き換える。
data_root = 'ADEChallengeData2016'
train 用スクリプトの用意
configs/segformer/segformer_mit-b0_8xb2-160k_ade20k-512x512.py を configs/segformer/segformer_ade.py にコピーする。
$ cp configs/segformer/segformer_mit-b0_8xb2-160k_ade20k-512x512.py configs/segformer/segformer_ade.py
segformer_ade.py スクリプト中の dataset のパスを変更する。
_base_ = [
'../_base_/models/segformer_mit-b0.py', '../_base_/datasets/ade20k_local.py', # ade20k_local.py に変更
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
]
コピー元のスクリプトを変更すると、使用モデル（ネットワークのサイズ）を変更できます。b0 が最も小さいモデルで学習が早く b5 が最も大きなモデルで学習に時間がかかります。b0 はチェックポイントのファイルサイズが 12Mb ほど、b5 は 300MB ほどです。
$ ls -1 configs/segformer/*.py
configs/segformer/segformer_mit-b0_8xb1-160k_cityscapes-1024x1024.py
configs/segformer/segformer_mit-b0_8xb2-160k_ade20k-512x512.py
configs/segformer/segformer_mit-b1_8xb1-160k_cityscapes-1024x1024.py
configs/segformer/segformer_mit-b1_8xb2-160k_ade20k-512x512.py
configs/segformer/segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py
configs/segformer/segformer_mit-b2_8xb2-160k_ade20k-512x512.py
configs/segformer/segformer_mit-b3_8xb1-160k_cityscapes-1024x1024.py
configs/segformer/segformer_mit-b3_8xb2-160k_ade20k-512x512.py
configs/segformer/segformer_mit-b4_8xb1-160k_cityscapes-1024x1024.py
configs/segformer/segformer_mit-b4_8xb2-160k_ade20k-512x512.py
configs/segformer/segformer_mit-b5_8xb1-160k_cityscapes-1024x1024.py
configs/segformer/segformer_mit-b5_8xb2-160k_ade20k-512x512.py
configs/segformer/segformer_mit-b5_8xb2-160k_ade20k-640x640.py
b0 以外のファイルは b0 を内部で参照しているため、データセットを変更する場合は、最初に上に示した segformer_ade.py を作成した上で、このファイルを参照する必要があります。例えば segformer_mit-b5_8xb2-160k_ade20k-512x512.py をベースにしたい場合は、segformer_ade_b5.py などにコピーした上で、下記のように冒頭で参照するファイルを書き換えます。
_base_ = ['./segformer_ade.py']
学習を実行
$ python tools/train.py configs/segformer/segformer_ade.py
epochs などすべてデフォルトにして b0 モデル使用、RTX3070Ti で4時間ほどかかりました。b5 の場合は 14時間ほどかかりました。
+---------------------+-------+-------+
| Class | IoU | Acc |
+---------------------+-------+-------+
| wall | 65.89 | 80.51 |
| building | 76.79 | 87.42 |
| sky | 92.48 | 95.67 |
| floor | 70.82 | 82.01 |
| tree | 67.43 | 84.88 |
| ceiling | 75.49 | 85.51 |
| road | 74.9 | 82.58 |
| bed | 72.91 | 93.12 |
| windowpane | 48.47 | 64.32 |
| grass | 60.14 | 78.03 |
| cabinet | 44.59 | 65.23 |
| sidewalk | 52.77 | 75.26 |
| person | 67.23 | 84.98 |
| earth | 32.19 | 43.97 |
| door | 25.12 | 34.1 |
| table | 35.32 | 50.99 |
| mountain | 50.2 | 69.82 |
| plant | 42.39 | 52.82 |
| curtain | 55.65 | 75.25 |
| chair | 34.89 | 51.57 |
| car | 73.28 | 86.9 |
| water | 35.5 | 47.58 |
| painting | 54.35 | 74.07 |
| sofa | 44.45 | 66.97 |
| shelf | 31.01 | 47.09 |
| house | 43.07 | 69.0 |
| sea | 50.64 | 82.51 |
| mirror | 28.48 | 35.53 |
| rug | 43.94 | 58.57 |
| field | 22.24 | 43.52 |
| armchair | 16.54 | 25.08 |
| seat | 46.15 | 59.01 |
| fence | 29.15 | 40.8 |
| desk | 20.03 | 34.13 |
| rock | 37.63 | 76.03 |
| wardrobe | 39.74 | 63.83 |
| lamp | 39.46 | 54.06 |
| bathtub | 42.15 | 62.48 |
| railing | 25.58 | 34.17 |
| cushion | 30.34 | 39.23 |
| base | 10.79 | 16.1 |
| box | 10.87 | 13.32 |
| column | 33.39 | 44.15 |
| signboard | 23.23 | 32.92 |
| chest of drawers | 27.24 | 40.74 |
| counter | 25.93 | 32.28 |
| sand | 35.19 | 42.36 |
| sink | 34.83 | 42.56 |
| skyscraper | 52.95 | 74.8 |
| fireplace | 40.03 | 72.49 |
| refrigerator | 33.32 | 57.38 |
| grandstand | 19.73 | 62.79 |
| path | 20.15 | 29.56 |
| stairs | 28.94 | 32.36 |
| runway | 69.13 | 93.7 |
| case | 39.71 | 67.21 |
| pool table | 79.12 | 95.27 |
| pillow | 28.48 | 36.19 |
| screen door | 33.52 | 75.52 |
| stairway | 25.46 | 39.03 |
| river | 7.74 | 18.76 |
| bridge | 35.88 | 60.79 |
| bookcase | 28.94 | 52.55 |
| blind | 17.4 | 19.98 |
| coffee table | 37.6 | 62.69 |
| toilet | 59.26 | 79.79 |
| flower | 21.13 | 35.51 |
| book | 41.27 | 55.25 |
| hill | 5.43 | 6.78 |
| bench | 30.64 | 44.38 |
| countertop | 36.22 | 53.39 |
| stove | 37.99 | 48.21 |
| palm | 37.03 | 56.32 |
| kitchen island | 17.23 | 52.68 |
| computer | 36.34 | 64.66 |
| swivel chair | 27.71 | 47.0 |
| boat | 42.77 | 72.16 |
| bar | 15.61 | 17.73 |
| arcade machine | 19.32 | 39.09 |
| hovel | 23.11 | 47.38 |
| bus | 57.59 | 81.89 |
| towel | 29.07 | 38.71 |
| light | 31.51 | 33.79 |
| truck | 3.39 | 5.35 |
| tower | 25.05 | 29.54 |
| chandelier | 45.16 | 67.53 |
| awning | 10.31 | 12.75 |
| streetlight | 7.61 | 8.72 |
| booth | 20.47 | 22.98 |
| television receiver | 40.89 | 61.34 |
| airplane | 34.87 | 61.25 |
| dirt track | 0.0 | 0.0 |
| apparel | 9.12 | 14.28 |
| pole | 7.21 | 8.56 |
| land | 5.69 | 6.26 |
| bannister | 0.51 | 0.54 |
| escalator | 23.21 | 32.21 |
| ottoman | 9.61 | 10.3 |
| bottle | 1.74 | 1.94 |
| buffet | 30.91 | 42.47 |
| poster | 0.76 | 0.76 |
| stage | 4.58 | 12.4 |
| van | 7.0 | 7.76 |
| ship | 32.94 | 41.58 |
| fountain | 9.93 | 11.46 |
| conveyer belt | 43.13 | 78.06 |
| canopy | 9.27 | 13.51 |
| washer | 47.19 | 67.22 |
| plaything | 4.66 | 7.09 |
| swimming pool | 28.62 | 79.52 |
| stool | 7.55 | 8.44 |
| barrel | 11.14 | 79.02 |
| basket | 8.98 | 10.31 |
| waterfall | 47.95 | 60.75 |
| tent | 60.37 | 99.52 |
| bag | 0.0 | 0.0 |
| minibike | 47.34 | 64.33 |
| cradle | 47.11 | 74.68 |
| oven | 8.38 | 13.59 |
| ball | 35.64 | 53.59 |
| food | 47.79 | 55.96 |
| step | 0.0 | 0.0 |
| tank | 24.63 | 25.04 |
| trade name | 15.99 | 17.47 |
| microwave | 20.69 | 24.11 |
| pot | 24.03 | 27.53 |
| animal | 39.2 | 39.92 |
| bicycle | 36.04 | 67.34 |
| lake | 0.41 | 0.49 |
| dishwasher | 16.54 | 31.76 |
| screen | 54.78 | 78.82 |
| blanket | 0.0 | 0.0 |
| sculpture | 24.36 | 27.34 |
| hood | 9.85 | 18.05 |
| sconce | 8.24 | 8.76 |
| vase | 16.07 | 23.59 |
| traffic light | 11.03 | 15.45 |
| tray | 0.45 | 0.5 |
| ashcan | 12.33 | 16.18 |
| fan | 30.89 | 37.08 |
| pier | 19.75 | 67.39 |
| crt screen | 0.0 | 0.0 |
| plate | 34.43 | 46.29 |
| monitor | 0.82 | 0.9 |
| bulletin board | 26.98 | 35.3 |
| shower | 0.0 | 0.0 |
| radiator | 26.01 | 27.87 |
| glass | 1.16 | 1.17 |
| clock | 8.43 | 8.84 |
| flag | 10.8 | 11.1 |
+---------------------+-------+-------+
11/30 01:11:32 - mmengine - INFO - Iter(val) [2000/2000] aAcc: 74.3100 mIoU: 30.1900 mAcc: 43.4200 data_time: 0.0012 time: 0.0111
最小のモデル(b0)で学習させた場合、上のような感じでした。
チェックポイントやログなど、デフォルトでは work_dir/segformer_ade 以下に出力されます。
$ ls work_dirs/segformer_ade/
20241129_205709 iter_112000.pth iter_160000.pth iter_48000.pth iter_96000.pth
20241129_205737 iter_128000.pth iter_16000.pth iter_64000.pth last_checkpoint
20241129_205808 iter_144000.pth iter_32000.pth iter_80000.pth segformer_ade.py
推論
demo/demo.py を demo/seg_demo.py にコピーして下記のように改造します。
from mmseg.apis import inference_model, init_model, show_result_pyplot
config_file = 'configs/segformer/segformer_ade.py' # config を相対パスで指定
checkpoint_file = 'work_dirs/segformer_ade/iter_160000.pth' # チェックポイントを相対パスで指定
model = init_model(config_file, checkpoint_file, device='cuda:0')
img = 'demo/demo.png' # or img = mmcv.imread(img), which will only load it once
result = inference_model(model, img) # segmentation したいファイルを img で指定
show_result_pyplot(model, img, result, show=False, out_file='result_segformer.jpg', opacity=0.5) # show=True とすると画面上に表示される、out_file に出力ファイルを指定
上の例では、チェックポイントは最終のものを指定しています。
$ python demo/seg_demo.py
b0 での出力結果はこんなかんじです。
やや微妙ですが、一応 segmentation できてますね。
b5 では上記のようになりました。b0 と比べるとセグメントの境界がややクリアになって、認識できる物体が増えています。
ちなみに元画像は下記です。
demo/image_demo.py にあるように、動画を対象として推論もできます。