获取VGG16特征层输出

import torch
import torchvision
from PIL import Image
from matplotlib import pyplot as plt
import torchvision.transforms as transforms
from torchvision.models import vgg16
from model_fcn8s import VGG

pretrain_backbone = True
struct = [(2, 64), (2, 128), (3, 256), (3, 512), (3, 512)]
backbone = VGG(num_classes=21, struct=struct)
if pretrain_backbone is True:
    backbone.load_state_dict(torch.load("/home/hyq/hyq/projects/fcn/vgg16_bn.pth"), strict=False)
 
return_layers = {'layer3':"layer3",'layer4':'layer4','layer5':"layer5"}
# 我的model，提取特征 vgg.features
backbone = torchvision.models._utils.IntermediateLayerGetter(backbone,return_layers)


imgPath = "/home/hyq/hyq/projects/fcn/test_01.jpg"
input= transforms.ToTensor()(Image.open(imgPath)).unsqueeze(0)


x = input
out = [x.squeeze(0).detach().numpy().transpose(2, 1, 0)[..., 0:3]]
#	给每个子图加上title
titles = ["pool 1", "pool 2", "pool 3", "pool 4", "pool 5"]
#	创建了一个具有5个子图的图像
fig, axs = plt.subplots(1, 5, figsize=(15, 3))

# for index, layer in enumerate(backbone.values()):
#     print(index, layer, type(backbone))

for index, (layer_name, layer) in enumerate(backbone.items()):
    print(index, layer_name, layer, type(backbone))
    
    # 打印所有池化层的输出
    if index in [0, 1, 2, 3, 4]:
        x = layer(x)
        print(f"Output shape after layer {index + 1}: {x.shape}")
        out.append(x.squeeze(0).detach().numpy().transpose(2, 1, 0)[..., 0:3])
        # plt.imshow(out[-1])
        # plt.show()
        # plt.savefig('output.png')
        axs[index].imshow(out[-1])
        axs[index].axis('off')
        axs[index].set_title(titles[index])
plt.tight_layout()
plt.show()
plt.savefig('output.png')