はじめに

mmsegmentationとはsemantic segmentationを実現するためのpytorchベースのToolboxです。
多数のアーキテクチャのモデルが提供されており、config一つで学習・推論できます。
この記事ではmmsegmentationのモデルで中間層の特徴量を取り出すコードを紹介します。

環境準備

ライブラリのインストール

git clone https://github.com/open-mmlab/mmsegmentation.git
cd mmsegmentation
pip install -v -e .
pip install mmcv-full

動作確認環境
- python 3.6.8
- mmsegmentation(0.29.1)
  - configを読み込む場合、githubからcloneした方が良いかもしれません。
  - wgetだとconfigがうまく開けませんでした。
- mmcv-full(1.7.0)

モデルのダウンロード

checkpoint_fileは使用したいモデルに応じて変更してください。
基本的にはmmsegmentationの各モデルのREADMEにcheckpointのパスが記載されています。
今回はこちらを参照しました。
https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3plus

checkpoint_file = 'https://download.openmmlab.com/mmsegmentation/v0.5/deeplabv3plus/deeplabv3plus_r50-d8_512x1024_40k_cityscapes/deeplabv3plus_r50-d8_512x1024_40k_cityscapes_20200605_094610-d222ffcd.pth'

!pip install wget
# google colabで実行する場合は”wget $checkpoint_file”でOK
!python -m wget $checkpoint_file

中間層の出力を得る方法

pytorchのmodelはnn.Moduleクラスを継承しています。
nn.Moduleクラスのregister_forward_hookメソッドにforward後に呼び出す関数を引数指定することで、中間層の出力を得ることが出来ます。
他にも中間層の出力を得る方法はあるようですが、今回はregister_forward_hookを使用していきます。

コード

指定した中間層の出力を得る関数

def middle_layer_output(target, inputs):
    """
    中間層の特徴量を得る関数
    """
    feature = None

    # forward後に呼び出される関数
    def forward_hook(module, inputs, outputs):
        global features
        # detach()で計算グラフから切り離し、clone()で値の共有をしないようにする
        features = outputs.detach().clone()

    handle = target.register_forward_hook(forward_hook)
    model.eval()
    # NOTE: mmsegmentationでない場合はmodel(inputs)で良い
    inference_segmentor(model, inputs)
    handle.remove()
    return features

こちらの関数を使用します。

コード全体

from pathlib import Path
import mmcv
from mmseg.apis import init_segmentor, inference_segmentor

mmseg_path = Path('./mmsegmentation') # path to your mmsegmentation
config_file_path = mmseg_path / 'configs/deeplabv3plus/deeplabv3plus_r50-d8_512x1024_40k_cityscapes.py'
checkpoint_file_path = './deeplabv3plus_r50-d8_512x1024_40k_cityscapes_20200605_094610-d222ffcd.pth'

# configの読み込み
cfg = mmcv.Config.fromfile(config_file_path)

# 複数GPUでBNを行うSyncBNが定義されており、single CPU/GPUで実行するとエラーになってしまうため、configのSyncBNをBNに書き換える
cfg.norm_cfg = dict(type='BN', requires_grad=True)
cfg.model.backbone.norm_cfg = cfg.norm_cfg
cfg.model.decode_head.norm_cfg = cfg.norm_cfg
cfg.model.auxiliary_head.norm_cfg = cfg.norm_cfg

# モデルの読み込み
model = init_segmentor(cfg, checkpoint_file_path, device='cpu')

def middle_layer_output(target, inputs):
    """
    中間層の特徴量を得る関数
    """
    feature = None

    # forward後に呼び出される関数
    def forward_hook(module, inputs, outputs):
        global features
        # detach()で計算グラフから切り離し、clone()で値の共有をしないようにする
        features = outputs.detach().clone()

    handle = target.register_forward_hook(forward_hook)
    model.eval()
    # NOTE: mmsegmentationでない場合はmodel(inputs)で良い
    inference_segmentor(model, inputs)
    handle.remove()
    return features

# 特徴量を得たい中間層を指定
target = model.backbone.layer4[-1].conv3
img = 'demo.jpg'
middle_layer_output(target, img)

References