如何正确使用AdaptiveAvgPool实现全局池化？解决输出维度问题

问题原因 #

AdaptiveAvgPool1d(1)输出的形状是[64, 32, 1]，而PyTorch的Linear层默认只对最后一个维度做线性变换，因此会把[64,32,1]转换成[64,32,4]，和你期望的[64,4]维度不符，最终导致和标签（形状[64,]）广播运算时出现维度不匹配错误。

对比原模型：Flatten会把[64,32,252]直接展平为[64, 8064]，Linear层基于这个一维特征输出[64,4]，能正常匹配标签维度。

修复方案 #

在AdaptiveAvgPool1d之后添加Flatten（或nn.Squeeze()），把[64,32,1]压缩成[64,32]，这样Linear层就会基于32维的全局池化特征输出[64,4]，和原模型输出维度完全一致。

修改后的模型代码 #

model = nn.Sequential(
        nn.Conv1d(in_channels=1, out_channels=32, kernel_size=128, stride=16, padding=1),
        nn.ReLU(),
        nn.MaxPool1d(kernel_size=2, stride=2),
        nn.AdaptiveAvgPool1d(1),
        nn.Flatten(),  # 新增该层压缩维度
        nn.LazyLinear(n_classes)
)

修改后的模型结构输出 #

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
Sequential                               --                        --
├─Conv1d: 1-1                            [64, 32, 505]             4,128
├─ReLU: 1-2                              [64, 32, 505]             --
├─MaxPool1d: 1-3                         [64, 32, 252]             --
├─AdaptiveAvgPool1d: 1-4                 [64, 32, 1]               --
├─Flatten: 1-5                           [64, 32]                  --
├─Linear: 1-6                            [64, 4]                   132
==========================================================================================

修改后Linear层输出形状为[64,4]，可正常与标签[64,]进行运算，解决维度不匹配的报错。

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