TensorFlow实现VGG16遇输入维度不兼容错误求助

问题描述 #

使用TensorFlow实现VGG16进行图像分类，通过tensorflow.data.Dataset创建的train_data的element_spec为：

<TensorSliceDataset element_spec=(TensorSpec(shape=(215, 160, 3), dtype=tf.float64, name=None), TensorSpec(shape=(20,), dtype=tf.float64, name=None))>

自定义了CustomConv2d层和VGG16模型类，初始化模型时用了model = VGG16((215,160,3,1))，拟合数据时出现错误：

ValueError: Input 0 of layer "conv11" is incompatible with the layer: expected min_ndim=4, found ndim=3. Full shape received: (215, 160, 3)

解决方案 #

1. 修正模型初始化的输入形状 #

input_shape参数仅需指定单张图像的形状（高、宽、通道数），不需要额外添加冗余维度，也无需包含批次维度。将模型初始化代码改为：

model = VGG16((215, 160, 3))

2. 确保数据集输入为4维张量 #

TensorFlow的Conv2D层要求输入是4维张量：(批次大小, 图像高度, 图像宽度, 通道数)，而当前数据集输出的是单张3维图像。需要在数据管道中添加批处理操作：

# 替换32为你需要的批次大小
train_data = train_data.batch(batch_size=32)

执行后train_data的element_spec会变为：

<TensorSliceDataset element_spec=(TensorSpec(shape=(None, 215, 160, 3), dtype=tf.float64, name=None), TensorSpec(shape=(None, 20,), dtype=tf.float64, name=None))>
其中None代表批次大小可变。

3. 可选优化 #

• 你的VGG16模型中conv_11使用了53个滤波器，标准VGG16第一个卷积层是64个滤波器，若需要对齐标准结构可修改为64。
• 模型定义中声明了self.dropout但未使用，若需要防止过拟合，可在全连接层之间添加dropout操作：

  def call(self,x,training=True):
      x = self.maxpool(self.conv_12(self.conv_11(x)))
      x = self.maxpool(self.conv22(self.conv21(x)))
      x = self.maxpool(self.conv33(self.conv32(self.conv31(x))))
      x = self.maxpool(self.conv43(self.conv42(self.conv41(x))))
      x = self.maxpool(self.conv53(self.conv52(self.conv51(x))))
      
      x = self.flatten(x)
      x = self.dropout(x, training=training)
      x = self.dense1(x)
      x = self.dropout(x, training=training)
      x = self.dense2(x)
      x = self.dense3(x)
      return x 
  

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