python-3.x - 查找 torchvision 模型中的所有 ReLU 层

Question

从 torchvision.models 获取预训练模型后，我希望将所有 ReLU 实例注册到 register_backward_hook(f)，其中是这样的:

for pos, module in self.model.features._modules.items():
    for sub_module in module:
        if isinstance(module, ReLU):
            module.register_backward_hook(f)

对我来说，问题是如何找到模型中的所有 ReLU 。对于densenet161，ReLU不仅存在于model.features._modules中，还存在于自定义的密集层中，例如: model.features._modules['denseblock1'][0]。对于resnet151，ReLU存在于model._modules及其自定义层中，例如model._modules['layer1'] .

有没有办法找到模型内的所有ReLU？

Answer 1

迭代模型所有组件的更优雅的方法是使用 modules()方法:

from torch import nn

for module in self.model.modules():
  if isinstance(module, nn.ReLU):
    module.register_backward_hook(f)

如果您不想获取所有子模块，而只想获取直接子模块，则可以考虑使用 children()方法而不是 modules()。您还可以使用 named_modules() 获取子模块的名称。方法。