Pytorch提取不同层次图片的特征

下例使用torchvision库提取了resnet最后一层的卷积特征；resnet各block的卷积特性，以及金字塔特性。

具体取哪一层特征视使用场景而定，resnet各block的输出包含更丰富的特征；从resnet最后一层提取的特征更为抽象；fpn每层通道数相等，含义也类似，可以在多层之间比较。

import os
import torchvision.models.detection.backbone_utils as backbone_utils
import torchvision
import torch.nn as nn

device = 'cuda'
os.environ["TORCH_HOME"] = '/notebooks/data/mine/live/code_v7/model/'
USE_FPN = True
if USE_FPN:
    backbone = backbone_utils.resnet_fpn_backbone('resnet50', True)
    features = list(backbone.children())[:-1] # 去掉最后的fpn层, 得到resnet的2,3,4层输出
    #features = list(backbone.children()) # 计算图像金字塔输出, 低层包括具体和抽像特征
    model = nn.Sequential(*features)   
else:
    backbone = torchvision.models.resnet50(pretrained=True)
    features = list(backbone.children())[:-2] # 去掉全连接和池化层, 得到最后卷积层输出 
    model = nn.Sequential(*features)
    
model = model.to(device)
x = torch.rand([1,3,244,244]).to(device)
out = model(x)

if USE_FPN: # 多层输出
    for key,value in out.items():
        print(key, value.shape)
else: # 单层输出
    print(out.shape)