PyTorch 中是否有内置函数可以打印 CNN 每个神经元的输出

Question

我创建了一个具有 4 个卷积层的 CNN 模型。如何找到每一层的神经元数量和输出

[CNN架构]

()```

Answer 1

不，没有内置函数可以输出 CNN 中每层的大小。

但是，有一些方法可以使用前向钩子检查模块：

nn.Module.register_forward_hook

。一种方法是在推理之前准备一个空字典，并通过钩子回调提供结果。

这是一个使用构造钩子工厂函数的最小示例：

get_parameter_count = lambda x: sum(p.numel() for p in x.parameters())

def create_hook(name, out):
    def hook(module, in_tensor, out_tensor):
        out[name] = get_parameter_count(module), out_tensor
    return hook

现在您可以将自定义挂钩附加到名为

model

的模块的不同层上：

out = {}
for name, layer in model.named_children():
    layer.register_forward_hook(create_hook(name, out))

最后我们准备通过调用

model

进行推断。中间结果和子模块大小将包含在

out

推理后。

---

这是包含以下类的演示：

class DummyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.convs = nn.Sequential(nn.Conv2d(3, 16, 3, 3), 
                                   nn.Conv2d(16, 8, 3, 3))
        self.pool = nn.AdaptiveAvgPool2d(2)
        self.classif = nn.Linear(32, 10)

    def forward(self, x):
        out = self.convs(x)
        out = self.pool(out)
        out = out.flatten(1)
        out = self.classif(out)
        return out

所需的字典由以下内容构建：

>>> model = DummyNet()
>>> y = model(torch.rand(2, 3, 100, 100))

>>> for name, (c, y) in out.items():
...     print(f'[{name}] #parameter={c}, output shape={y[0].shape}')
[convs] #parameter=1608, output shape=(2, 8, 11, 11)
[pool] #parameter=0, output shape=(2, 8, 2, 2)
[classif] #parameter=330, output shape=(2, 10)

您可以在另一个线程中阅读有关前向钩子的更多信息：如何获取图中的所有张量？