时间戳

首页 > 解决方案 > 如何在 PyTorch 模型中查找中间层的输入层名称？

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

问题描述

我在 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

标签： pythonpytorchlayer

解决方案

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

推荐阅读