PyTorch中卷积层连接全连接层时如何推导输出形状确定输入维度

你不需要手动逐层记录输出张量形状来计算X，你当前使用的(W - F + 2P) / S + 1逐层计算公式是准确的，适合需要明确每一层输出尺寸的网络设计场景，同时也有更便捷的自动计算方法可以使用，能避免手动计算的失误。

最推荐的自动计算方法：Dummy张量前向推导 #

直接构造一个和实际输入形状完全一致的随机张量，传入卷积层后将输出展平，得到的维度就是X的取值，示例代码如下：

import torch
import torch.nn as nn

# 你的卷积层定义
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))

# 替换下方参数为你实际的输入尺寸：[batch_size, 通道数, 图片高度, 图片宽度]，batch_size填1即可
dummy_input = torch.randn(1, 3, 32, 32)
conv_output = conv_layer(dummy_input)
# 展平除batch维度外的所有维度，得到的长度就是X
X = conv_output.flatten(start_dim=1).shape[1]
print(f"第一个全连接层输入维度X为：{X}")

后续如果修改了卷积层的参数（比如调整卷积核大小、新增padding、增减卷积/池化层），只要重新运行这段代码就能拿到最新的X值，不需要重新逐层计算。

可选自动方案：模型结构总结工具 #

你也可以使用第三方的模型结构打印工具，直接输出每一层的输出形状：

• 用pip安装pytorch-model-summary工具
• 导入后调用summary方法，传入模型和输入尺寸，就能直接看到所有层的输出维度，卷积层最后一层的输出展平后就是X的取值。

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