【pytorch学习】hook打印中间特征图、计算网络算力等

r囧r小猫 2022-11-05 01:56 253阅读 0赞

> 点击“**机器学习与生成对抗网络**”，关注星标
> 
> 获取有趣、好玩的前沿干货！

作者：知乎—涩醉

地址：https://www.zhihu.com/people/se-zui

**01**

**背景**

在神经网络的反向传播当中个，流程只保存叶子节点的梯度，对于中间变量的梯度没有进行保存。

import torchx = torch.tensor([1,2],dtype=torch.float32,requires_grad=True)y = x+2z = torch.mean(torch.pow(y, 2))lr = 1e-3z.backward()x.data -= lr*x.grad.dataprint(y.grad)

此时输出就是：None，这个时候hook的作用就派上，hook可以通过自定义一些函数，从而完成中间变量的输出，比如中间特征图、中间层梯度修正等。

在pytorch docs搜索hook，可以发现有四个hook相关的函数，分别为register\_hook，register\_backward\_hook，register\_forward\_hook，register\_forward\_pre\_hook。其中register\_hook属于tensor类，而后面三个属于moudule类。

*  register\_hook函数属于torch.tensor类，函数在tensor梯度计算的时候就会执行，这个函数主要处理梯度相关的数据，表现形式  .

import torchx = torch.tensor([1,2],dtype=torch.float32,requires_grad=True)y = x * 2y.register_hook(print)<torch.utils.hooks.RemovableHandle at 0x7f765e876f60>z = torch.mean(y)z.backward()tensor([ 0.5000,  0.5000])

*  Register\_backward\_hook等三个属于torch.nn，属于moudule中的方法。

hook(module, grad_input, grad_output) -> Tensor or None

写个demo，参考(https://oldpan.me/archives/pytorch-autograd-hook):

下面的计算为

![format_png][]

![format_png 1][]

![format_png 2][]

import torchimport torch.nn as nndevice = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    def print_hook(grad):    print ("register hook:", grad)    return grad
    class TestNet(nn.Module):    def __init__(self):        super(TestNet, self).__init__()        self.f1 = nn.Linear(4, 1, bias=True)        self.weights_init()
        def weights_init(self):        self.f1.weight.data.fill_(4)        self.f1.bias.data.fill_(0.1)
        def forward(self, input):        self.input = input        out = input * 0.75        out = self.f1(out)        out = out / 4        return out
        def back_hook(self, moudle, grad_input, grad_output):        print ("back hook in:", grad_input)        print ("back hook out:", grad_output)        # 修改梯度        # grad_input = list(grad_input)        # grad_input[0] = grad_input[0] * 100        # print (grad_input)        return tuple(grad_input)
    
    if __name__ == '__main__':    input = torch.tensor([1, 2, 3, 4], dtype=torch.float32, requires_grad=True).to(device)    net = TestNet()    net.to(device)    net.register_backward_hook(net.back_hook)    ret = net(input)    print ("result", ret)    ret.backward()    print('input.grad:', input.grad)    for param in net.parameters():        print('{}:grad->{}'.format(param, param.grad))

输出：

result tensor([7.5250], grad_fn=<DivBackward0>)back hook in: (tensor([0.2500]), None)back hook out: (tensor([1.]),)input.grad: tensor([0.7500, 0.7500, 0.7500, 0.7500])Parameter containing:tensor([[4., 4., 4., 4.]], requires_grad=True):grad->tensor([[0.1875, 0.3750, 0.5625, 0.7500]])Parameter containing:tensor([0.1000], requires_grad=True):grad->tensor([0.2500])

输出结果以及梯度都很明显，简单分析一下w权重的梯度，

![format_png 3][]

![format_png 4][]

另外，hook中有个bug，假设我们bug，假设我们注释掉out = out / 4这行，可以发现输出变成back hook in: (tensor(\[1.\]), tensor(\[1.\]))。这种情况就不符合上面我们的梯度计算公式，是因为这个时候：

![format_png 5][]

则此时的偏导只是对 ![format_png 6][] 和 ![format_png 7][] 进行计算，所以都是1，1。这是pytorch的设计缺陷

![4a766a9193d33bc5120e122dd063b57b.png][]

*  register\_forward\_hook跟Register\_backward\_hook差不多，就不过多复述。  
    
 *  register\_forward\_pre\_hook，可以发现其输入只有  
    hook(module, input) -> None  
    其主要是针对推理时的hook.

###     ###

**02**

**应用**

### **2.1 特征图打印** ###

直接利用pytorch已有的resnet18进行特征图打印，只打印卷积层的特征图，

import torchfrom torchvision.models import resnet18import torch.nn as nnfrom torchvision import transforms
    import matplotlib.pyplot as plt
    def viz(module, input):    x = input[0][0]    #最多显示4张图    min_num = np.minimum(4, x.size()[0])    for i in range(min_num):        plt.subplot(1, 4, i+1)        plt.imshow(x[i].cpu())    plt.show()
    
    import cv2import numpy as npdef main():    t = transforms.Compose([transforms.ToPILImage(),                            transforms.Resize((224, 224)),                            transforms.ToTensor(),                            transforms.Normalize(mean=[0.485, 0.456, 0.406],                                                 std=[0.229, 0.224, 0.225])                            ])
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        model = resnet18(pretrained=True).to(device)    for name, m in model.named_modules():        # if not isinstance(m, torch.nn.ModuleList) and \        #         not isinstance(m, torch.nn.Sequential) and \        #         type(m) in torch.nn.__dict__.values():        # 这里只对卷积层的feature map进行显示        if isinstance(m, torch.nn.Conv2d):            m.register_forward_pre_hook(viz)    img = cv2.imread('./cat.jpeg')    img = t(img).unsqueeze(0).to(device)    with torch.no_grad():        model(img)
    if __name__ == '__main__':    main()

直接放几张中间层的图

![24cb15dee151d4e4c74914cf53e66985.png][]

图1 第一层卷积层输入

![0ba5c381602419fbef472495aaf30c66.png][]

图2 第四层卷积层的输入

### **2.2 模型大小，算力计算** ###

同样的用法，可以直接参考pytorch-summary(https://github.com/sksq96/pytorch-summary)这个项目。

**03**

**参考**

https://oldpan.me/archives/pytorch-autograd-hook

https://pytorch.org/docs/stable/search.html?q=hook&check\_keywords=yes&area=default

https://github.com/pytorch/pytorch/issues/598

https://github.com/sksq96/pytorch-summary

https://github.com/allensll/test/blob/591c7ce3671dbd9687b3e84e1628492f24116dd9/net\_analysis/viz\_lenet.py

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

**猜您喜欢：**

[**超100篇！CVPR 2020最全GAN论文梳理汇总！**][100_CVPR 2020_GAN]

[拆解组新的GAN：解耦表征MixNMatch][GAN_MixNMatch]

[StarGAN第2版：多域多样性图像生成][StarGAN_2]

[附下载 | 《可解释的机器学习》中文版][Link 1]

[附下载 |《TensorFlow 2.0 深度学习算法实战》][_TensorFlow 2.0]

[附下载 |《计算机视觉中的数学方法》分享][Link 2]

[《基于深度学习的表面缺陷检测方法综述》][Link 3]

[《零样本图像分类综述: 十年进展》][Link 4]

[《基于深度神经网络的少样本学习综述》][Link 5]

![9205ad7556569239766a1b54c20cd139.png][]

[format_png]: /images/20221023/bd9461061231446d86a0bdc11905d3c7.png
[format_png 1]: /images/20221023/bb9de5cb4b5346fd842782d9d61bdddd.png
[format_png 2]: /images/20221023/ef4b91f68bda4458acf49b6e51170c4c.png
[format_png 3]: /images/20221023/eb34db88e5fc4791943a9d8135d0c450.png
[format_png 4]: /images/20221023/9b590b20dc514e7bb0e4c559b521f2d0.png
[format_png 5]: /images/20221023/8be13942203941b4afd8c6f5ae222418.png
[format_png 6]: https://imgconvert.csdnimg.cn/aHR0cHM6Ly9tbWJpei5xcGljLmNuL21tYml6X3N2Zy9XbXdxanNTQnNaSXFZTDkxUG4wS1NjR1JXOThtbGNnQXRkWmNoY21scW5rNjJtVUdWNDZ1dEF3SG03SVJJOWZYMUI1N0wwaWNDaWMyRjYxUDkyTGxsN2ljOFBSWGdrQkdnWkovNjQw?x-oss-process=image/format,png
[format_png 7]: /images/20221023/8ebe7d10ef7e40798062eb38eb7d7eb2.png
[4a766a9193d33bc5120e122dd063b57b.png]: /images/20221023/8acffa71c18542b5b2561832ea1f9a06.png
[24cb15dee151d4e4c74914cf53e66985.png]: /images/20221023/455bb7c5ee88491394797729ce014c18.png
[0ba5c381602419fbef472495aaf30c66.png]: /images/20221023/843f271d30ba401781d5de053aea0823.png
[100_CVPR 2020_GAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a94cc95e205ae880cb9518cb16a54c0abde69e5a8654f7fedf70a0cc990775265bf28a1c&idx=1&mid=2247486455&scene=21&sn=94156f53845cdf07e859e343862be745#wechat_redirect
[GAN_MixNMatch]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aaa9c95e23bf3131249a8d008ce3acdff7c52c7f6be84d92c6d607ca804b5ceb2a4222e3&idx=1&mid=2247485458&scene=21&sn=9f562c1de84dada67268f88c72ff77ff#wechat_redirect
[StarGAN_2]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29ab6ec95e22785abb2afbb2e2b1c95780a4d6bfceeb915118440b41a81769a8776dcc62b7&idx=1&mid=2247485909&scene=21&sn=048d1336b8083bac3755e27a1f968d06#wechat_redirect
[Link 1]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a4bac95e2dac758c4aac10a3aac49eb062bc0829d4ef2ebc8e6d5069d246385c5084a2b4&idx=1&mid=2247484929&scene=21&sn=523220858287ed3e6105919703c8e833#wechat_redirect
[_TensorFlow 2.0]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a5b0c95e2ca6599d4a91d5d40900566b7525de3dd85259e7a6c173aa26aa7ea3641d1aaf&idx=2&mid=2247485195&scene=21&sn=e3387067db29d1759755c6c2ea96c870#wechat_redirect
[Link 2]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa8fc95e23993e057fb0dbe09f2810ea07184a439c173f404a7ab657cadcde3d24efa52a&idx=2&mid=2247485492&scene=21&sn=c4fb135de02cd52f06bc61b0e79476e4#wechat_redirect
[Link 3]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa8fc95e2399a93dae3af3b13626643b18cf42842d1a7ada0de46486fa8b865aa2421d2f&idx=1&mid=2247485492&scene=21&sn=6f8af433521b5fee5aceab62e0111115#wechat_redirect
[Link 4]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a44fc95e2d597fc88ca7f10cae079f821deedb6f93b10f0985aa8f6f7c06e369d716f785&idx=1&mid=2247485172&scene=21&sn=f65b58ed3b4c1d09163271069dce4738#wechat_redirect
[Link 5]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a7e6c95e2ef0a6649c0c904fa28b61b4ac76774df4acdce8fff3a909369c7c2f8d22a4b5&idx=1&mid=2247484765&scene=21&sn=b1b9a9ef6c923abad6cf38364105264c#wechat_redirect
[9205ad7556569239766a1b54c20cd139.png]: /images/20221023/3c14858f93364f2dbe51b88b7ba6b25d.png