python - 如何在 PyTorch 模型中查找中间层的输入层名称？

Question

我在 PyTorch 上有一些复杂的模型。如何打印连接到图层输入的图层名称(或 ID)。首先，我想为 Concat 层找到它。请参阅下面的示例代码:

class Concat(nn.Module):
    def __init__(self, dimension=1):
        super().__init__()
        self.d = dimension

    def forward(self, x):
        return torch.cat(x, self.d)


class SomeModel(nn.Module):
    def __init__(self):
        super(SomeModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        self.conc = Concat(1)
        self.linear = nn.Linear(8192, 1)

    def forward(self, x):
        out1 = F.relu(self.bn1(self.conv1(x)))
        out2 = F.relu(self.conv2(x))
        out = self.conc([out1, out2])
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out


if __name__ == '__main__':
    model = SomeModel()
    print(model)
    y = model(torch.randn(1, 3, 32, 32))
    print(y.size())
    for name, m in model.named_modules():
        if 'Concat' in m.__class__.__name__:
            print(name, m, m.__class__.__name__)
            # Here print names of all input layers for Concat

Answer 1

您可以使用 type(module).__name__ 获取 nn.Module 类名:

>>> model = SomeModel()
>>> y = model(torch.randn(1, 3, 32, 32))
>>> for name, m in model.named_modules():
...     if 'Concat' == type(m).__name__:
...         print(name, m)
conc Concat()

---

编辑:您实际上可以设法获取用于计算 Concat 输入的运算符列表。但是，我认为您实际上无法获取与这些运算符关联的 nn.Module 的属性名称。此类信息在模型推理时不可用且不需要。

此解决方案要求您使用 nn.Module.register_forward_hook 在层上注册前向 Hook .然后执行一个推理触发它，然后你可以删除钩子(Hook)。在前向 Hook 中，您可以访问输入列表并从 grad_fn 属性回调中提取运算符的名称。在这里使用 nn.Module.register_forward_pre_hook 会更合适，因为我们只查看输入，不需要输出。

>>> def op_name(x)
...     return type(x.grad_fn).__name__.replace('Backward0', '')

>>> def forward_hook(module, ins):
...     print([op_name(x) for x in ins[0]])

在 model.conc 上附加钩子(Hook)，触发它然后清理:

>>> handle = model.conc.register_forward_pre_hook(forward_hook)
>>> model(torch.empty(2, 3, 10, 10, requires_grad=True))
['Relu', 'Relu']

>>> handle.remove()