当卷积层与密集层连接时，如何推断输出的形状?

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]推导出来的。

Answer 2

如果你不想遍历各层并迭代地计算输出形状，你可以通过离线定义模型的CNN部分来进行推理:

cnn = 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))

例如，如果你有一个形状为(1, 3, 100, 100):

的输入

>>> cnn(torch.empty(1, 3, 100, 100)).shape
torch.Size([1, 128, 19, 19])

所以第一个完全连接的致密层的特征数应该是:

>>> cnn(torch.empty(1, 3, 100, 100)).numel()
46208