MMSegmentation の SegFormer で学習と推論をする (Ubuntu 22.04)

# やること

特にひねりはなくタイトルの通りですが、OpenMMLab のページの ReadMe に書かれている通りにやるとできないので、できない部分を補足しています。SegFormer 以外を試す場合でもインストール周りはたぶん同じです。

2024/12現在、最低でも半年以上は ReadMe が修正されないままになってます

# 前提

- Ubuntu 22.04
- Cuda 12.2
- NVIDIA Driver Version 535.183.01 (Cuda のバージョンに合わせる）

cuda 12.5 などにする場合は、driver のバージョンをそれに合わせる必要があります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/266247/cf8578f3-986e-1874-9a3a-0c6a134430e0.png)

cuda と driver のバージョンの合わせ方については、NVIDIA のダウンロードのページにいってインストーラを選択したときに、コマンドラインでダウンロードする時の `deb` ファイルの名前を見ます。

```bash
$ 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 のインストール

```bash
$ conda create -n mmseg$
$ conda activate mmseg
$ conda install python=3.8 -y
$ conda install pytorch torchvision -c pytorch
```

python は 3.8 の使用が絶対です（多分）。

## Pytorch で cuda が使用可能かテスト

```bash
~$ 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 のインストール

```bash
$ 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 のバージョンが合わない場合などにもビルドエラーになります。

```bash
$ git clone https://github.com/open-mmlab/mmsegmentation.git
$ cd mmsegmentation
$ pip install -v -e .
```

以上でインストールは終了

# 学習テスト

## データの用意

ADEChallengeData2016.zip を使う場合、データは HuggingFace や Kaggle から入手可能。

- [HuggingFace](https://huggingface.co/hngan/my_lasagna/tree/b1ea139bf7a03c7c0a84af238dd27239bcf73345)
- [Kaggle](https://www.kaggle.com/datasets/ipythonx/ade20k-scene-parsing)

展開すると下記のようになります。ここでは `mmsegmentation` フォルダ直下に配置するものとします。

```bash
ADEChallengeData2016/
├── annotations
│   ├── training
│   └── validation
└── images
    ├── training
    └── validation
```

## dataset のコンフィグの用意

`configs/_base_/datasets/ade20k.py` を `configs/_base_/datasets/ade20k_local.py` にコピーする。

```bash 
$ cp configs/_base_/datasets/ade20k.py configs/_base_/datasets/ade20k_local.py
```

`configs/_base_/datasets/ade20k_local.py` の `data_root` をデータを配置したパスに書き換える。

```configs/_base_/datasets/ade20k_local.py
data_root = 'ADEChallengeData2016'
```

## train 用スクリプトの用意

`configs/segformer/segformer_mit-b0_8xb2-160k_ade20k-512x512.py` を `configs/segformer/segformer_ade.py` にコピーする。

```bash
$ cp configs/segformer/segformer_mit-b0_8xb2-160k_ade20k-512x512.py configs/segformer/segformer_ade.py
```

`segformer_ade.py` スクリプト中の `dataset` のパスを変更する。

```configs/segformer/segformer_ade.py
_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 ほどです。

```bash
$ 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` などにコピーした上で、下記のように冒頭で参照するファイルを書き換えます。

```segformer_ade_b5.py
_base_ = ['./segformer_ade.py']
```

## 学習を実行

```bash
$ python tools/train.py configs/segformer/segformer_ade.py
```

epochs などすべてデフォルトにして `b0` モデル使用、RTX3070Ti で4時間ほどかかりました。`b5` の場合は 14時間ほどかかりました。


```bash
+---------------------+-------+-------+
|        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` にコピーして下記のように改造します。

```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 に出力ファイルを指定
```

上の例では、チェックポイントは最終のものを指定しています。

```bash
$ python demo/seg_demo.py
```

`b0` での出力結果はこんなかんじです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/266247/336409bc-e856-68d6-7d69-8637d2d88096.png)

やや微妙ですが、一応 segmentation できてますね。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/266247/71f0a070-6405-c9b1-bb9c-344bb5b89425.png)

`b5` では上記のようになりました。`b0` と比べるとセグメントの境界がややクリアになって、認識できる物体が増えています。

ちなみに元画像は下記です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/266247/767c1880-79bb-c48e-664d-c9d25eb9cd43.png)

- 参照元: https://github.com/open-mmlab/mmsegmentation/blob/main/demo/demo.png

`demo/image_demo.py` にあるように、動画を対象として推論もできます。