卷积层特征提取可视化

写在前面，代码白嫖时间长久，忘记是在哪里白嫖的了，emmmmm，如果有人发现出处，麻烦评论一下，我看见了就引用。
代码在colab平台测试可用，效果图也放在下面，只能保证代码可以用。

方法1

import matplotlib.pyplot as plt
import torch
from torchvision import models, transforms
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import os
import numpy as np
from torchvision.datasets import ImageFolder

!unzip -d /content/ck /content/CK_Pro.zip

torch.cuda.set_device(0)  # 设置GPU ID
is_cuda = True
simple_transform = transforms.Compose([transforms.Resize((224, 224)),
              transforms.ToTensor(),  # H, W, C -> C, W, H 归一化到(0,1)，简单直接除以255
              transforms.Normalize([0.485, 0.456, 0.406],  # std
              [0.229, 0.224, 0.225])])

# mean  先将输入归一化到(0,1)，再使用公式”(x-mean)/std”，将每个元素分布到(-1,1)
# 使用 ImageFolder 必须有对应的目录结构
train = ImageFolder("/content/ck/CK_Pro/test", simple_transform)
# valid = ImageFolder("./datas/dogs-vs-cats/valid", simple_transform)
train_loader = DataLoader(train, batch_size=1, shuffle=False)
# val_loader = DataLoader(valid, batch_size=1, shuffle=False, num_workers=5)

vgg = models.vgg16(pretrained=True).cuda()

# 提取不同层输出的 主要代码
class LayerActivations:
    features = None

    def __init__(self, model, layer_num):
        self.hook = model[layer_num].register_forward_hook(self.hook_fn)

    def hook_fn(self, module, input, output):
        self.features = output.cpu()

    def remove(self):
        self.hook.remove()


print(vgg.features)

conv_out = LayerActivations(vgg.features, 0)  # 提出第 一个卷积层的输出
img = next(iter(train_loader))[0]

# imshow(img)
o = vgg(Variable(img.cuda()))
conv_out.remove()  #
act = conv_out.features  # act 即 第0层输出的特征

# 可视化 输出
fig = plt.figure()
# fig.subplots_adjust(left=0, right=0.5, bottom=0, top=0.8, hspace=0, wspace=0)
for i in range(64):
    ax = fig.add_subplot(8, 8, i + 1, xticks=[], yticks=[])
    ax.imshow(act[0][i].detach().numpy(), cmap="gray")

plt.show()

方法2

代码没跑出来，emmm，不想改代码