pytorch模型中间层特征的提取

转载自：https://blog.csdn.net/LXX516/article/details/80132228

定义一个特征提取的类：

 #中间特征提取  class FeatureExtractor(nn.Module):      def __init__(self, submodule, extracted_layers):          super(FeatureExtractor,self).__init__()          self.submodule = submodule          self.extracted_layers= extracted_layers         def forward(self, x):          outputs = []          for name, module in self.submodule._modules.items():              if name is "fc": x = x.view(x.size(0), -1)              x = module(x)              print(name)              if name in self.extracted_layers:                  outputs.append(x)          return outputs

 #特征输出  myresnet=resnet18(pretrained=False)  myresnet.load_state_dict(torch.load('cafir_resnet18_1.pkl'))   exact_list=["conv1","layer1","avgpool"]  myexactor=FeatureExtractor(myresnet,exact_list)  x=myexactor(img)

在这里主要应用的是:

for nama, module in model._modules.items():

所以要根据自己的情况重写这个类，这个类提供个一个很不错的想法

文章来源: pytorch模型中间层特征的提取