特征提取和可视化

如果是自己构建的模型，那么可以再forward函数中，返回特定层的输出特征图。如果是预训练模型，建议使用create_feature_extractor获取指定层的输出的方法。

针对自己构建的模型

创建一个可以根据需要调用的特征提取器类，代码如下所示：

class Feature_extractor(nn.module):
   def forward(self, input):
      self.feature = input.clone()
      return input
new_net = nn.Sequential().cuda() # the new network
target_layers = [conv_1, conv_2, conv_4] # layers you want to extract`
i = 1
for layer in list(cnn):
   if isinstance(layer,nn.Conv2d):
      name = "conv_"+str(i)
      art_net.add_module(name,layer)
      if name in target_layers:
         new_net.add_module("extractor_"+str(i),Feature_extractor())
      i+=1
   if isinstance(layer,nn.ReLU):
      name = "relu_"+str(i)
      new_net.add_module(name,layer)
   if isinstance(layer,nn.MaxPool2d):
      name = "pool_"+str(i)
      new_net.add_module(name,layer)
new_net.forward(your_image)

print (new_net.extractor_3.feature)

预训练模型

import torchvision
from PIL import Image
import torchvision.transforms as transforms
from matplotlib import pyplot as plt
from torchvision.models.feature_extraction import create_feature_extractor

transform = transforms.Compose([transforms.Resize(256),
                                transforms.CenterCrop(224),
                                transforms.ToTensor(),
                                transforms.Normalize([0.485, 
                                0.456, 0.406], [0.229, 0.224, 0.225])])

model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)

feature_extractor = create_feature_extractor(model, return_nodes={"conv1": "output"})

original_img = Image.open("dog.jpg")

img = transform(original_img).unsqueeze(0)

out = feature_extractor(img)

plt.imshow(out["output"][0].transpose(0, 1).sum(1).detach().numpy())

plt.show()