将卷积层与密集层连接时如何推断输出的形状？

Question

我正在尝试使用构建卷积神经网络pytorch，但无法理解如何解释第一个密集连接层的输入神经元。举例来说，我有以下架构：

self.conv_layer = nn.Sequential(
   nn.Conv2d(3, 32, 5),
   nn.Conv2d(32, 64, 5),
   nn.MaxPool2d(2, 2),
   nn.Conv2d(64, 128, 5),
   nn.Conv2d(128, 128, 5),
   nn.MaxPool2d(2, 2))

self.fc_layer = nn.Sequential(
   nn.Linear(X, 512),
   nn.Linear(512, 128),
   nn.Linear(128, 10))

这X是第一个线性层中的神经元数量。那么，我是否需要跟踪每一层输出张量的形状以便我可以弄清楚X？

现在，我可以将值放入公式中(W - F + 2P) / S + 1并计算每层之后的形状，这会有点方便。

难道没有更方便的东西可以自动执行此操作吗？

Answer 1

一个简单的解决方案是使用LazyLinear层： https: //pytorch.org/docs/stable/ generated/torch.nn.LazyLinear.html 。

根据文档：

torch.nn.Linear推断出的模块...in_features它们将在第一次调用完成后初始化forward，并且该模块将成为常规torch.nn.Linear模块。in_features线性的论证是从 推断出来的input.shape[-1]。