如何在 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

如何打印连接到图层输入的图层（或 ID）的名称。

您不能基于模块本身来执行此操作，因为“连接到”不作为 nn.Module 的静态属性存在：

代替

out = self.conc([out1, out2])

你可以很容易地写出来

out = self.conc([out1, random.choice([out1, out2]))

现在与self.conc 有什么关联？

但是，您可以检查 值的“图表”（如果 torchviz.make_dot 没有直接执行您想要的操作，您可能需要检查其实现）

Answer 2

您可以使用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 注册一个前向挂钩。然后执行一个推理来触发它，然后你可以移除钩子。在 forward 钩子中，您可以访问输入列表并从 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上附加钩子，触发它然后清理：

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

>>> handle.remove()