如何在PyTorch中自动计算卷积层输出维度（无需nn.LazyLinear）？

我希望能自动计算任意PyTorch卷积层的输出维度。在构建大型CNN时，手动计算最后一层卷积的输出会非常繁琐。请问有没有无需使用nn.LazyLinear的实时自动计算方法？

代码示例 #

self.conv1 = nn.Conv2d(1, 16, kernel_size=5, stride=2)
self.conv2 = nn.Conv2d(16, 32, kernel_size=5, stride=2)
self.conv3 = nn.Conv2d(32, 32, kernel_size=5, stride=2)

self.feature_dim = [ADD CODE HERE]

self.fc1 = nn.Linear(self.feature_dim, 256)

我曾尝试使用公式[(W−K+2P)/S]+1并在代码中实现，但为每一层都这么做比手动计算还要麻烦。有没有统一的方法直接从PyTorch读取输出形状？

可行方案 #

可以通过以下两个函数在前向传播过程中提取任意阶段的输出形状，其中torch.numel()可返回给定张量的总元素数：

self.flat_features = self._get_conv_output((1, 84, 84))

def _get_conv_output(self, shape):
    with torch.no_grad():
        input = torch.rand(1, *shape)
        output = self._forward_features(input)
        return output.numel()

def _forward_features(self, x):
    x = F.relu(self.conv1(x))
    x = F.relu(self.conv2(x))
    x = F.relu(self.conv3(x))
    return x

内容的提问来源于stack exchange，提问作者Tristan B