CNN关键点检测报错：Can not squeeze dim[2], expected a dimension of 1, got 2

错误原因分析 #

你遇到的错误源于模型输出维度与标签（关键点）维度不匹配：

• 目标是预测4个(x,y)关键点，每个样本的标签形状应为(4,2)（对应4个点，每个点2个坐标），批量输入时标签形状是(batch_size,4,2)。
• 当前模型最后一层是Dense(4)，输出形状为(batch_size,4)，与标签的(batch_size,4,2)维度不兼容。MSE损失函数计算时尝试挤压维度，但发现dim[2]的维度是2而非1，因此报错。

修复方案 #

需要让模型输出维度与标签维度对齐，有两种可行方式：

方式1：直接输出扁平化的8个坐标值 #

将最后一层全连接层的神经元数量改为8（4个点×2个坐标），此时模型输出形状为(batch_size,8)，对应标签可扁平化处理为(batch_size,8)（比如将[[x1,y1],[x2,y2],...]转为[x1,y1,x2,y2,...]）。

方式2：输出与标签一致的(4,2)形状 #

保持最后一层输出8个神经元，再通过Reshape层将输出重塑为(4,2)，直接匹配标签的二维结构。

修改后的代码示例 #

以下是采用方式2的完整修改代码：

dataset = dataset.batch(batch_size)

# 可验证标签形状：取消注释查看keypoints维度
# for image,keypoints in dataset.as_numpy_iterator():
#     print("Image shape:",image.shape)
#     print("Keypoints shape:",keypoints.shape)

input_shape = (1000,1000,1)

model = Sequential()
model.add(Conv2D(3, (3,3), strides = (1,1), padding='same', activation='relu', input_shape=input_shape))
model.add(BatchNormalization())
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Flatten())
model.add(Dense(10,activation='relu'))
model.add(Dense(8))  # 4个点×2个坐标=8个输出
model.add(Reshape((4, 2)))  # 重塑为与标签一致的(4,2)形状
model.summary()

model.compile(loss='mse',optimizer='adam',metrics='mae')

print("Fit on Dataset")
history = model.fit(dataset, epochs=1, batch_size=batch_size)

额外验证建议 #

取消代码中标签形状的打印注释，确认keypoints.shape是否为(5,4,2)（对应batch_size=5），确保模型输出与标签维度完全匹配。

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