在 torchvision 模型中查找所有 ReLU 层答案

【问题标题】：Find all ReLU layer in a torchvision model在 torchvision 模型中查找所有 ReLU 层
【发布时间】：2019-03-09 06:49:44
【问题描述】：

从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？

【问题讨论】：

标签： python-3.x machine-learning deep-learning computer-vision pytorch

【解决方案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()方法获取子模块的名称。

【讨论】：