当需要获取网络中间层输出时,如果模型是自定义的,可以在模型定义时直接将想获取的那一层作为输出返回,如分类网络同时返回最终输出类别和最后一层特征:
但当我们调用封装好的模型,又不想重写模型,但还要获取网络中间输出时,就可使用hook机制
首先打印模型,查看待输出层的名称
print(model)
定义hook
取特征:
feature = fmap_block[‘input’]
取梯度:
grad = grad_block[‘grad_in’]
参考:
https://www.jianshu.com/p/69e57e3526b3