pytorch提取网络任意层输出的特征图

网上很多教程教我们怎么获取某一层网络输出的feature map。
原理很简单，对于我们自己写的网络其实很容易，我们只要在return 的时候返回我们想要的特征就可以了，但是对于一些别人写好的，我们又不方便改原代码的网络，提取特定层的feature map就有点麻烦了。
下面我以torchvision的vgg16为例，提取第一个全连接层输出的feature map
首先我们看看网络都由什么层

import torchvision
model = torchvision.models.vgg16(pretrained=True)
print(model)

输出：

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

可以看到这里有conv，ReLU，MaxPool2d，等等，我们让我们的输入依次经过这些层，然后获取我们需要的输出就可以了。

import torchvision
model = torchvision.models.vgg16(pretrained=True)
#实际实用的时候，a可以换成我们需要输出的图片
a=torch.zeros((2,3,224,224))
#我们枚举整个网络的所有层
for i,m in enumerate(model.modules()):
	#让网络依次经过和原先结构相同的层，我们就可以获取想要的层的输出
	if isinstance(m, nn.Conv2d) or isinstance(m, nn.BatchNorm2d) or\
			isinstance(m, nn.ReLU) or isinstance(m, nn.MaxPool2d) or isinstance(m, nn.AdaptiveAvgPool2d):
		print(m)
		a= m(a)
	#我只想要第一个全连接层的输出
	elif isinstance(m, nn.Linear):
		print(m)	
		#和源代码一样，将他展平成一维的向量
		a = torch.flatten(a, 1)
		#获取第一个全连接层的输出
		a= m(a)
		break
print(a)

这样子，a就是我们最后想要的第一个全连接层的输出了。

更新
没想到这么多人看，那再更新一些比较简单的方法。
当我们需要用一些像VGG之类的作为backbone时，只想要feature层，最后的全连接可以去掉，可以这样做：

import torchvision.models as m
import torch.nn as nn
model = m.resnet50(pretrained=True)
#直接将最后一层赋值为空,模型的输出就是池化后的特征图了
model.fc = nn.Sequential([])
input_ = torch.zeros((2,3,224,224))
output_ = model(input_)
#(2,2048)
print(output_.shape)