pytorch使用hook获得特征图

最新推荐文章于 2023-09-20 14:36:52 发布

Shuai@

最新推荐文章于 2023-09-20 14:36:52 发布

阅读量711

点赞数

分类专栏： pytorch 文章标签： pytorch 深度学习

本文链接：https://blog.csdn.net/weixin_37707670/article/details/125209597

版权

pytorch 专栏收录该内容

23 篇文章 2 订阅

订阅专栏

本文介绍了如何通过定义forward和backward hook函数来捕获卷积神经网络（LeNet）中的特征图和梯度信息。通过在conv2层注册forward_hook，输入随机张量并执行前向传播，可以得到不同层的特征图大小和形状。这种方法对于网络可视化和理解深度学习模型内部工作原理非常有用。

摘要由CSDN通过智能技术生成

更细节的建议参考使用hook获得特征图可视化

定义如下两个hook函数，简单几步就能获得特征图了


fmap_block = dict()  # 装feature map
grad_block = dict()  # 装梯度

def backward_hook(module, grad_in, grad_out):
    grad_block['grad_in'] = grad_in
    grad_block['grad_out'] = grad_out


def farward_hook(module, inp, outp):
    fmap_block['input'] = inp
    fmap_block['output'] = outp

import torch
from torch import nn

#定义网路
class LeNet(nn.Module):  
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3,6,3,1,1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(6,9,3,1,1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2,2)
        self.fc1 = nn.Linear(8*8*9, 120)
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(120,10)

    def forward(self, x):
        out = self.pool1(self.relu1(self.conv1(x)))
        out = self.pool2(self.relu2(self.conv2(out)))
        out = out.view(out.shape[0], -1)
        out = self.relu3(self.fc1(out))
        out = self.fc2(out)

        return out

#实例化网络
net = LeNet()


# 注册hook
net.conv2.register_forward_hook(farward_hook)

定义输入

input_img = torch.randn(1,3,32,32).requires_grad_()  
fmap_block = dict()  # 装feature map
outs = net(input_img)

print(len(fmap_block))
print(fmap_block["output"].shape)