Pytorch获取中间层输出

清纯世纪

已于 2022-07-04 08:58:41 修改

阅读量1.2k

点赞数

分类专栏：记录本文章标签： pytorch 深度学习人工智能

于 2022-06-30 23:34:18 首次发布

本文链接：https://blog.csdn.net/qq_45100200/article/details/125550167

版权

本文介绍了如何在PyTorch中使用register_hook()函数来获取深度学习模型的中间层输出。通过为特定模块设置回调函数，可以在前向传播过程中截取所需层的数据。在实践中，建议先确定层名以精确提取，代码示例展示了如何添加钩子以及如何处理多层输出的情况。注意，此方法每次模型预测都会保存一次数据，因此要获取所有样本的输出，需要逐个进行预测。

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

钩子截流：

这种方式是在前向传播进行中还没得到最终输出时，将所需要的中间层输出从前向数据流中提取出来，利用到了pytorch中的register_hook()函数。这一函数可以为模型中的某个module设置一个回调函数，形如：
hook(module, input, output) -> None or modified output
函数的输入值为module的名字、module的输入和输出。通过前置定义一个数组，在hook()函数中将对应module的输入或输出加入该数组以实现中间层提取。实际过程中建议先打印所有层的名字以做到精确提取。给出代码如下：

网络：


class net1(nn.Module):
    def __init__(self):
        super(net1, self).__init__()
 
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 3, kernel_size=3, stride=1, padding=0, bias=False),
            nn.Conv2d(3, 6, kernel_size=3, stride=1, padding=0, bias=False),
        )
        self.conv2 = nn.Sequential(
            nn.C