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

Question

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

[CNN Architecture]

(https://i.stack.imgur.com/XASaf.png) ' ' '

Answer 1

对你的问题给出一个简短的回答:No.

---

你可以通过附加钩子来创建模块检查器。一个简单的方法是在推理之前准备一个空字典，并通过钩子调用提供结果。

下面是一个使用nn.Module.register_forward_hook的最小示例。

构造钩子工厂函数:

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.LazyLinear(10)
def forward(self, x):
out = self.convs(x)
out = self.pool(out)
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=torch.Size([8, 11, 11])
[pool] #parameter=0, output shape=torch.Size([8, 2, 2])
[classif] #parameter=30, output shape=torch.Size([8, 2, 10])

你可以在另一个线程中阅读更多关于向前钩子的内容:如何在一个图中获得所有张量?