pytorch 如何在YOLOv7推理过程中提取任意图层的特征值并可视化?

dtcbnfnu  于 2023-02-04  发布在  其他
关注(0)|答案(1)|浏览(571)

在我的例子中,我希望在cfg/training/yolov7.yaml中的以下代码中提取并可视化层102、103、104中的特征输出。

# yolov7 head
head:
  [[-1, 1, SPPCSPC, [512]], # 51
  
   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [37, 1, Conv, [256, 1, 1]], # route backbone P4
   [[-1, -2], 1, Concat, [1]],
   
   [-1, 1, Conv, [256, 1, 1]],
   [-2, 1, Conv, [256, 1, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]],
   [-1, 1, Conv, [256, 1, 1]], # 63
   
   [-1, 1, Conv, [128, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [24, 1, Conv, [128, 1, 1]], # route backbone P3
   [[-1, -2], 1, Concat, [1]],
   
   [-1, 1, Conv, [128, 1, 1]],
   [-2, 1, Conv, [128, 1, 1]],
   [-1, 1, Conv, [64, 3, 1]],
   [-1, 1, Conv, [64, 3, 1]],
   [-1, 1, Conv, [64, 3, 1]],
   [-1, 1, Conv, [64, 3, 1]],
   [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]],
   [-1, 1, Conv, [128, 1, 1]], # 75
      
   [-1, 1, MP, []],
   [-1, 1, Conv, [128, 1, 1]],
   [-3, 1, Conv, [128, 1, 1]],
   [-1, 1, Conv, [128, 3, 2]],
   [[-1, -3, 63], 1, Concat, [1]],
   
   [-1, 1, Conv, [256, 1, 1]],
   [-2, 1, Conv, [256, 1, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]],
   [-1, 1, Conv, [256, 1, 1]], # 88
      
   [-1, 1, MP, []],
   [-1, 1, Conv, [256, 1, 1]],
   [-3, 1, Conv, [256, 1, 1]],
   [-1, 1, Conv, [256, 3, 2]],
   [[-1, -3, 51], 1, Concat, [1]],
   
   [-1, 1, Conv, [512, 1, 1]],
   [-2, 1, Conv, [512, 1, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]],
   [-1, 1, Conv, [512, 1, 1]], # 101
   
   [75, 1, RepConv, [256, 3, 1]],   #extract
   [88, 1, RepConv, [512, 3, 1]],   #extract
   [101, 1, RepConv, [1024, 3, 1]], #extract

   [[102,103,104], 1, IDetect, [nc, anchors]],   # Detect(P3, P4, P5)
  ]

然而,我希望能够采取任何层的特点,如果可能的话,因为我可能需要层的特点,而不是这一个。
我该怎么做呢?
我尝试从models/yolo.py中的Model类中提取并可视化https://github.com/ultralytics/yolov5/issues/3089,但无法确定要编辑哪些代码以及如何编辑。我尝试对IDetect类执行相同操作,但也无法确定。

0yg35tkg

0yg35tkg1#

你可以注册一个forward hook到相关的层。根据pytorch documentation,“每次forward()计算完输出后都会调用这个hook。”
从本质上讲,forward hook函数修改了一个全局作用域的变量,这个变量在层forward调用终止后仍然存在,你可以将层forward调用的输出(通过forward hook函数)存储在这个变量中,然后你可以在以后引用它。
(明确地说,我相信注册forward钩子会隐式地将nn.module更改为全局作用域,这样在函数调用终止后,该值会持续存在。更多信息请参见pytorch文档。)
在任何情况下,前向钩子函数都需要以下函数签名:

hook(module, input, output) -> None or modified output

举个简单的例子:

def make_hook(key):
    def hook(model, input, output):
        intermediate_output[key] = output.detach()
    return hook

外部函数本身返回一个函数,因为register_module_forward_hook的输入是一个具有上述签名的函数。
然后我们可以添加forward钩子到任何模块:

model.<layer_name>.register_forward_hook(make_hook("example_key"))

因此,总的来说,您的代码将类似于:

def make_hook(key):
    def hook(model, input, output):
        intermediate_output[key] = output.detach()
    return hook

# define model
model = Yolo5() # I know this is wrong but you didn't include the actual model in your question so this is just an example

# register hook to as many layers as you want
model.conv4.register_forward_hook("conv4") # same here, I made these layer names up
model.maxpool8.register_forward_hook("maxpool8")

# dummy input
inp = torch.random.rand(1,3,1080,1920)

# forward pass
model(inp)

# reference intermediate_outputs
intermediate_outputs["conv4"] # should have the output from this layer stored as value

请注意,因为使用前向钩子向module“添加全局状态”,pytorch文档建议仅将此特性临时用于调试目的,而不是持久解决方案。对于长期解决方案,您可以修改主模型架构的forward通道,以将这些值存储为中间输出,并在结束时返回所有这些值。

相关问题