【検証】Pytorchにおける特徴量抽出器の挙動

はじめに

深層学習では特徴量に注目した分析・学習を行うことが多いです。Pytorchにおける特徴量抽出器の作成方法=中間層の出力を見る方法としては以下のような手法があります。
・torchvision.models.feature_extraction.create_feature_extractorを使う
・必要な層のみをnn.Sequentialでまとめる
・(register_forward_hookを用いる)
参考文献
https://qiita.com/kzkadc/items/bb1bd536da5a5f8f488a

本記事ではこの特徴量抽出器が学習・推論時にどのような挙動をするかについてまとめていきます。

print(torch.__version__)
print(torchvision.__version__)
2.1.0+cu118
0.16.0+cu118

内部のパラメータの挙動

学習前と学習後でパラメータの値を比較して、元のモデルと特徴量抽出器のパラメータが共有されているのかについて確認します。

以下のコードを実際に実行してみると、学習前も学習後も元のモデルと特徴量抽出器は同じパラメータ値を出力します。このことから特徴量抽出器は元のモデルをコピーしているのではなく、元のモデルの一部分を参照しているだけにすぎないことが分かります。(コピーをする場合はcopy.deepcopy等を挟むのが丸い？)

import torch
from torch import nn, optim
from torchvision import datasets, models, transforms
from torchvision.models.feature_extraction import create_feature_extractor

# データセット
trainset = datasets.CIFAR10(root='./data', train=True, transform=transforms.ToTensor(), download=True)
trainloader = torch.utils.data.DataLoader(trainset, batch_size = 2)

# モデル
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
# create_feature_extractorを使用する場合
feature_extractor = create_feature_extractor(model, {"avgpool": "feature"})
# 必要な層のみnn.Sequentialでまとめる場合
# feature_extractor = nn.Sequential(*list(model.children())[:-1]) 

# 学習前
print("before")
print(list(model.conv1.parameters())[0][0,0,0])
print(list(feature_extractor.conv1.parameters())[0][0,0,0])
# print(list(feature_extractor[0].parameters())[0][0,0,0])

# 適当に学習
optimizer = optim.Adam(params=model.parameters())
criterion = nn.CrossEntropyLoss()

images, labels = next(iter(trainloader))
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

# 学習後
print("after")
print(list(model.conv1.parameters())[0][0,0,0])
print(list(feature_extractor.conv1.parameters())[0][0,0,0])
# print(list(feature_extractor[0].parameters())[0][0,0,0])

before
tensor([-0.0104, -0.0061, -0.0018,  0.0748,  0.0566,  0.0171, -0.0127],
       grad_fn=<SelectBackward0>)
tensor([-0.0104, -0.0061, -0.0018,  0.0748,  0.0566,  0.0171, -0.0127],
       grad_fn=<SelectBackward0>)
after
tensor([-0.0114, -0.0071, -0.0028,  0.0758,  0.0576,  0.0181, -0.0117],
       grad_fn=<SelectBackward0>)
tensor([-0.0114, -0.0071, -0.0028,  0.0758,  0.0576,  0.0181, -0.0117],
       grad_fn=<SelectBackward0>)

訓練モード、推論モードの挙動

Pytorchではbatchnormやdropoutといった学習時と推論時で異なる挙動をするモジュールに対して、訓練モード(model.train)と推論モード(model.eval)を用いることで対応しています。特徴量抽出器を用いた場合、これはどのように動作するのかを以下のコードで検証します。

import torch
from torch import nn, optim
from torchvision import datasets, models, transforms
from torchvision.models.feature_extraction import create_feature_extractor

trainset = datasets.CIFAR10(root='./data', train=True, transform=transforms.ToTensor(), download=True)
trainloader = torch.utils.data.DataLoader(trainset, batch_size = 2)
images, labels = next(iter(trainloader))

model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
feature_extractor = create_feature_extractor(model, {"avgpool": "feature"})

model.train()
feature_extractor.train()

print(model.training)
print(model.conv1.training)
print(feature_extractor.training)
print(feature_extractor.conv1.training)

print(model(images))
print(feature_extractor(images)['feature'])

# True
# True
# True
# True
# tensor([[-0.5343,  3.5019,  1.7081,  ...,  0.3263,  2.4757, -2.2703],
#         [ 0.1216, -4.4287, -2.9703,  ..., -1.1383,  0.4308,  4.1788]],
#        grad_fn=<AddmmBackward0>)
# tensor([[[[1.4333]],
#          [[2.0614]],
#          [[2.0009]],
#          ...,

model.eval()
feature_extractor.eval()

print(model.training)
print(model.conv1.training)
print(feature_extractor.training)
print(feature_extractor.conv1.training)

print(model(images))
print(feature_extractor(images)['feature'])

# False
# False
# False
# False
# tensor([[ 0.4302,  2.2151, -0.2216,  ..., -0.5493, -0.0544,  1.9519],
#         [-1.2921, -1.4514, -0.8693,  ..., -1.4664, -0.6607,  2.1136]],
#        grad_fn=<AddmmBackward0>)
# tensor([[[[0.7870]],
#          [[0.0000]],
#          [[0.0000]],
#          ...,

元のモデルと特徴量抽出器を両方訓練モードに設定した場合全ての要素が訓練モードになり、推論モードに設定した場合全ての要素が推論モードになる。(当たり前)

model.train()
feature_extractor.eval()

print(model.training)
print(model.conv1.training)
print(feature_extractor.training)
print(feature_extractor.conv1.training)

print(model(images))
print(feature_extractor(images)['feature'])

# True
# <False>
# False
# False
# tensor([[ 0.4302,  2.2151, -0.2216,  ..., -0.5493, -0.0544,  1.9519],
#         [-1.2921, -1.4514, -0.8693,  ..., -1.4664, -0.6607,  2.1136]],
#        grad_fn=<AddmmBackward0>)
# tensor([[[[0.7870]],
#          [[0.0000]],
#          [[0.0000]],
#          ...,

元のモデルを訓練モードにした後に特徴量抽出器を推論モードにすると、元モデル内部のconv1が推論モードのままになり、元モデルの出力も推論モード時と同等になっていることが読み取れます。

このことから、以下のことが言えそうです。
.train()と.eval()はモデルと全ての構成要素のモードを変換するという処理が行われている。モデル実行時は、モデルのモードではなく構成要素のモードを参照している。

そのため上記のように、model.train()をしているのにモデル出力が推論モードのものになるという現象が起きるようです。これは学習時にも影響がありそうなので気を付けた方がいい問題な気がします。

おわりに

自分が何気なくfeature_extractor.eval()をしたらmodelの学習の挙動も変化したことから、色々気になったことを実験してまとめてみました。今まで特に意識していなかったことなので以後気を付けたいです。