本文首发自【】用户【西北小生_】的博客，转载请注明出处！
PyTorch之HOOK——获取神经网络特征和梯度的有效工具记录了PyTorch获取卷积神经网络特征图和梯度的方法，但由于举例简单，并不能直接应用于用其它构造方法构造的神经网络模型。为解决更具一般性的问题，现针对PyTorch自带的vgg16模型，记录获取其任一层输出特征图的方法，读者阅读完本文，可以自行应用于获取更多的模型（resnet, densenet,mobilenet等）特征图和梯度。
首先导入需要的包：
import torch
from torch import nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from torchvision import models
from torchvision.transforms.functional import normalize, resize, to_tensor, to_pil_image
本文以pytorch自带的vgg16模型为例，故需要从torchvision.models加载vgg16模型:
model = models.vgg16_bn(pretrained=True).eval()
我们打印一下vgg16_bn看一下它的结构：
In [3]: print(model)
VGG(
(features): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
(6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(7): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(8): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(9): ReLU(inplace=True)
(10): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(12): ReLU(inplace=True)
(13): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(14): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(15): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(16): ReLU(inplace=True)
(17): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(18): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(19): ReLU(inplace=True)
(20): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(21): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(22): ReLU(inplace=True)
(23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(24): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(25): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(26): ReLU(inplace=True)
(27): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(28): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(29): ReLU(inplace=True)
(30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(31): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(32): ReLU(inplace=True)
(33): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(35): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(36): ReLU(inplace=True)
(37): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(38): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(39): ReLU(inplace=True)
(40): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(41): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(42): ReLU(inplace=True)
(43): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
(classifier): Sequential(
(0): Linear(in_features=25088, out_features=4096, bias=True)
(1): ReLU(inplace=True)
(2): Dropout(p=0.5, inplace=False)
(3): Linear(in_features=4096, out_features=4096, bias=True)
(4): ReLU(inplace=True)
(5): Dropout(p=0.5, inplace=False)
(6): Linear(in_features=4096, out_features=1000, bias=True)
)
)
可以看出pytorch自带的vgg16_bn模型是由features, avgpool, classifier 三个children构成的，features和classifier又由Sequential()组成。
如果我们想获取features中任意一层的输出特征图或梯度，该如何获取呢？
那我们就取男人最爱的18这个数字吧（手动狗头）
假设要获取features中的编号为18的层输出的特征图，我们首先要建立hook函数，然后对该层用register_forward_hook()函数进行注册（对这一部分不熟悉的可以看我的博客）：
# 建立一个全局变量的字典，将特征图放在其中
feature_map = {}
# 构建hook函数
def forward_hook(module, inp, outp):
feature_map['features18'] = outp
其实这一方法的难点就在于如何获取相应的层，并对其注册hook，我介绍一种简单的做法，那就是torch.nn.Module.children()方法。（torch.nn.Module.children()介绍请看我的这篇博客）
pytorch自带的vgg16_bn模型有3个children：features, avgpool, classifier 。为了获取features中的第18层，我们需要先获取features，再用索引的方式获取第18层（或者任一层）：
features = list(model.children())[0]
hook_layer = features[18]
这样就获取了features中的第18层，在vgg16_bn中，第18层是BN层，我们可以打印hook_layer变量验证一下：
In [11]: hook_layer
Out[11]: BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
可以看到，我们确实已经获取到了第18层。只要看懂了这里，获取其它层都是一样简单的。
接下来就是对第18层进行hook注册：
hook_layer.register_forward_hook(forward_hook)
到这里，只需要加载一幅图像，输入模型model，对其进行前向传播一次，第18层的输出特征图就会出现在全局字典变量feature_map中了。
加载一幅图像（本例中的图像来源于ImageNet）：
imgdir = 'ILSVRC2012_val_00000003.JPEG'
origin_img = Image.open(imgdir)
img_tensor = normalize(to_tensor(resize(origin_img, (224, 224))),
[0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
input_img = torch.unsqueeze(img_tensor, 0)
将图像输入模型中，进行一次前向传播：
with torch.no_grad():
score = model(input_img)
这时第18层输出的特征图就已经出现在feature_map中了，我们查看一下它的尺寸：
In [23]: feature_map['features18'].shape
Out[23]: torch.Size([1, 256, 56, 56])
可以看到我们已经获取了第18层的输出特征图，后续处理就不一一赘述。等以后有时间了写一个CAM系列讲hook的应用。