Hook 函数与 CAM 算法( 二 ) _CAM

输出如下：
forward_pre_hook input:(tensor([[[[1., 1., 1., 1.],[1., 1., 1., 1.],[1., 1., 1., 1.],[1., 1., 1., 1.]]]]),)backward hook input:(None, tensor([[[[0.5000, 0.5000, 0.5000],[0.5000, 0.5000, 0.5000],[0.5000, 0.5000, 0.5000]]],[[[0.5000, 0.5000, 0.5000],[0.5000, 0.5000, 0.5000],[0.5000, 0.5000, 0.5000]]]]), tensor([0.5000, 0.5000]))backward hook output:(tensor([[[[0.5000, 0.0000],[0.0000, 0.0000]],[[0.5000, 0.0000],[0.0000, 0.0000]]]]),)hook函数实现机制hook函数实现的原理是在module的__call()__函数进行拦截，__call()__函数可以分为 4 个部分：

第 1 部分是实现 _forward_pre_hooks
第 2 部分是实现 forward 前向传播
第 3 部分是实现 _forward_hooks
第 4 部分是实现 _backward_hooks

由于卷积层也是一个module，因此可以记录_forward_hooks 。

def __call__(self, *input, **kwargs):# 第 1 部分是实现 _forward_pre_hooksfor hook in self._forward_pre_hooks.values():result = hook(self, input)if result is not None:if not isinstance(result, tuple):result = (result,)input = result# 第 2 部分是实现 forward 前向传播if torch._C._get_tracing_state():result = self._slow_forward(*input, **kwargs)else:result = self.forward(*input, **kwargs)# 第 3 部分是实现 _forward_hooksfor hook in self._forward_hooks.values():hook_result = hook(self, input, result)if hook_result is not None:result = hook_result# 第 4 部分是实现 _backward_hooksif len(self._backward_hooks) > 0:var = resultwhile not isinstance(var, torch.Tensor):if isinstance(var, dict):var = next((v for v in var.values() if isinstance(v, torch.Tensor)))else:var = var[0]grad_fn = var.grad_fnif grad_fn is not None:for hook in self._backward_hooks.values():wrapper = functools.partial(hook, self)functools.update_wrapper(wrapper, hook)grad_fn.register_hook(wrapper)return result

Hook 函数提取网络的特征图下面通过hook函数获取 AlexNet 每个卷积层的所有卷积核参数，以形状作为 key，value 对应该层多个卷积核的 list 。然后取出每层的第一个卷积核，形状是 [1, in_channle, h, w]，转换为 [in_channle, 1, h, w]，使用 TensorBoard 进行可视化，代码如下：

writer = SummaryWriter(comment='test_your_comment', filename_suffix="_test_your_filename_suffix")# 数据path_img = "imgs/lena.png"# your path to imagenormMean = [0.49139968, 0.48215827, 0.44653124]normStd = [0.24703233, 0.24348505, 0.26158768]norm_transform = transforms.Normalize(normMean, normStd)img_transforms = transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor(),norm_transform])img_pil = Image.open(path_img).convert('RGB')if img_transforms is not None:img_tensor = img_transforms(img_pil)img_tensor.unsqueeze_(0)# chw --> bchw# 模型alexnet = models.alexnet(pretrained=True)# 注册hookfmap_dict = dict()for name, sub_module in alexnet.named_modules():if isinstance(sub_module, nn.Conv2d):key_name = str(sub_module.weight.shape)fmap_dict.setdefault(key_name, list())# 由于AlexNet 使用 nn.Sequantial 包装，所以 name 的形式是：features.0features.1n1, n2 = name.split(".")def hook_func(m, i, o):key_name = str(m.weight.shape)fmap_dict[key_name].append(o)alexnet._modules[n1]._modules[n2].register_forward_hook(hook_func)# forwardoutput = alexnet(img_tensor)# add imagefor layer_name, fmap_list in fmap_dict.items():fmap = fmap_list[0]# 取出第一个卷积核的参数fmap.transpose_(0, 1) # 把 BCHW 转换为 CBHWnrow = int(np.sqrt(fmap.shape[0]))fmap_grid = vutils.make_grid(fmap, normalize=True, scale_each=True, nrow=nrow)writer.add_image('feature map in {}'.format(layer_name), fmap_grid, global_step=322)

使用 TensorBoard 进行可视化如下：