Pythonいぬ

pythonを使った画像処理に関する記事を書いていきます

Pytorchで全学習済みモデルの特徴マップと特徴ベクトルの抽出

転移学習、スタイル変換、物体検知、セマンティックセグメンテーション、メトリックラーニング、perceptual loss、ゼロショット学習など学習済みモデルの中間層を使いたい場合がよくある。Pytorchで使える学習済みモデルの特徴マップと特徴ベクトルを抽出する方法についてまとめてみる。

特徴マップと特徴ベクトル

図で書くと以下のようになる。

f:id:tzmi:20200501165032p:plain

CNNの出力で一次元ベクトルになる直前が特徴マップ(Feature map)で、最終層のひとつ手前が特徴ベクトル(Feature vector)。pytorchで使えるそれぞれの学習済みモデルでこれらの特徴を抽出してみる。

pytorchで使える学習済みモデル

まずはpytorchで使える学習済みモデルについて確認

In [1]: from torchvision import models

In [2]: dir(models)                                                             
Out[2]: 
['AlexNet',
 'DenseNet',
 'GoogLeNet',
 'GoogLeNetOutputs',
 'Inception3',
 'InceptionOutputs',
 'MNASNet',
 'MobileNetV2',
 'ResNet',
 'ShuffleNetV2',
 'SqueezeNet',
 'VGG',
:
:
:
]

使える学習済みネットワークがたくさんある。こんなに増えたのか。。。

学習済みモデルを調査する際にちょっとハマったんだけど、python3.8にすると、pipでインストールできるtorchvisionの最新バージョンは0.2.2になる。バージョン0.2.2にはmobilenetなどが入っていないので注意。

python3.7にすると最新版(20/01/25時点で0.5.0)がインストールできる。 minicondaなどで環境作るときは

conda create -n pthfeat python=3.7 conda

というように python=3.7 を入れて環境構築するように注意する。

AlexNet

特徴マップ

from torchvision import models
model = models.alexnet(pretrained=True).features

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 256, 6, 6])

特徴ベクトル

from torchvision import models
from torch import nn

model = models.alexnet(pretrained=True)
layers = list(model.classifier.children())[:-1]
model.classifier = nn.Sequential(*layers)

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 4096])

VGG

特徴マップ

from torchvision import models
model = models.vgg16(pretrained=True).features

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 512, 7, 7])

特徴ベクトル

from torchvision import models
from torch import nn

model = models.vgg16(pretrained=True)
layers = list(model.classifier.children())[:-2]
model.classifier = nn.Sequential(*layers)

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 4096])

ResNet, ResNext, WideResNet

特徴マップ

from torchvision import models
from torch import nn

model = models.resnet50(pretrained=True)
layers = list(model.children())[:-2]
model = nn.Sequential(*layers)

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 2048, 7, 7])

特徴ベクトル

from torchvision import models
from torch import nn

model = models.resnet50(pretrained=True)
model.fc = nn.Identity()

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 2048])

SqueezeNet

特徴マップ

from torchvision import models
model = models.squeezenet1_0(pretrained=True).features

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズは torch.Size([1, 512, 13, 13])

特徴ベクトル

squeezenetは最終層手前にLinearレイヤを持たないため、特徴ベクトルは特徴マップのサイズを1x1にしたものとなる

from torchvision import models
from torch import nn
model = models.squeezenet1_0(pretrained=True)
model.classifier = nn.AdaptiveAvgPool2d((1,1))

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズは torch.Size([1, 512])

DenseNet

特徴マップ

from torchvision import models
model = models.densenet121(pretrained=True).features

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズは torch.Size([1, 1024, 7, 7])

特徴ベクトル

from torchvision import models
model = models.densenet121(pretrained=True)
model.classifier = nn.Identity() # 恒等関数に変換

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズは torch.Size([1, 1024])

GoogLeNet

特徴マップ

from torchvision import models
from torch import nn

model = models.googlenet(pretrained=True)
layers = list(model.children())[:-3]
model = nn.Sequential(*layers)

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 1024, 7, 7])

特徴ベクトル

from torchvision import models
from torch import nn

model = models.googlenet(pretrained=True)
layers = list(model.children())[:-2]
model = nn.Sequential(*layers, nn.Flatten())

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 1024])

ShuffleNet

特徴マップ

from torchvision import models
from torch import nn

model = models.shufflenet_v2_x1_0(pretrained=True)
layers = list(model.children())[:-1]
model = nn.Sequential(*layers) 

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 1024, 7, 7])

特徴ベクトル

from torchvision import models
from torch import nn

model = models.shufflenet_v2_x1_0(pretrained=True)
model.fc = nn.Identity() # 恒等関数に変換

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズはtorch.Size([1, 1024])

MobileNet

特徴マップ

from torchvision import models
model = models.mobilenet_v2(pretrained=True).features

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズは torch.Size([1, 1280, 7, 7])

特徴ベクトル

from torchvision import models
from torch import nn

model = models.mobilenet_v2(pretrained=True)
model.classifier = nn.Identity() # 恒等関数に変換

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズは torch.Size([1, 1280])

MNASNet

特徴マップ

from torchvision import models

model = models.mnasnet1_0(pretrained=True).layers

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズは torch.Size([1, 1280, 7, 7])

特徴ベクトル

from torchvision import models
from torch import nn

model = models.mnasnet1_0(pretrained=True)
model.classifier = nn.Identity() # 恒等関数に変換

入力サイズがtorch.Size([1, 3, 224, 224])に対して、出力サイズは torch.Size([1, 1280])

まとめていて気になったこと

各ネットワークともに複雑な構造はモジュール化されているので、ほとんどがSequentialを使ってもよかったが、Global Average Pooling 2d をやるときにAdaptiveAvgPoolingを使っているパターンとx.mean([2,3])を使っているパターンがあった。前者はモジュール扱いなのでSequentialに入るが、後者はモジュールではないため、Sequentialにすると無視される。なので、できれば最終層のモジュール(fcとかclassifierとか)をnn.Identity()を使って恒等関数のモジュールに変更する方法が一番確実なんじゃないかと思った。

参考サイト

github.com