TensorFlow输入944×944图像遇维度错误及训练损失NaN问题求助

咱们一步步来拆解你遇到的三个问题，逐个解决：

1. 先搞定Reshape维度不匹配的错误 #

你的第一个报错完全是因为Dense层前的Reshape操作写错了特征图尺寸，咱们手动算一下经过两次池化后的特征图大小：

• 输入图像是944×944，第一次卷积用了same padding，尺寸不变；接着用2×2的最大池化、步长2，尺寸减半成 472×472
• 第二次卷积同样是same padding，尺寸保持472×472；再一次2×2池化后，尺寸再次减半成 236×236

所以pool2的实际形状是[batch_size, 236, 236, 64]，你之前写的944*944*64完全不符合当前的特征图尺寸，这就导致Reshape时维度无法整除，直接触发错误。修改这一行代码即可：

# 替换原来的错误Reshape
pool2_flat = tf.reshape(pool2, [-1, 236*236*64])

2. 解决训练损失为NaN的问题 #

修复Reshape后出现NaN损失，核心原因是任务配置不匹配：你做的是二分类，但模型按多分类在配置，咱们分两种方案调整：

方案一：用Softmax实现二分类（简单适配原代码） #

• 首先把Logits层的输出单元数改成2（对应两个类别）：

  logits = tf.layers.dense(inputs=dropout, units=2)
  

• 损失函数可以继续用sparse_softmax_cross_entropy（因为你的标签是0/1的整数，刚好符合这个损失的输入要求）。如果依然出现NaN，可以把学习率从0.001调低到0.0001试试，避免权重更新幅度过大。

方案二：用Sigmoid实现二分类（更适合二分类任务） #

这是二分类任务更常用的方式，调整步骤如下：

1. 修改Logits层为单输出（对应二分类的概率）：

   logits = tf.layers.dense(inputs=dropout, units=1)
   

2. 调整损失函数时，必须保证标签和logits的形状一致：
   你的标签是[batch_size,]的一维数组，而logits是[batch_size,1]的二维数组，需要先把标签reshape成相同形状，再计算损失：

   # 转换标签类型并调整形状
   labels = tf.cast(tf.reshape(labels, [-1, 1]), tf.float32)
   # 计算交叉熵并取batch内的平均值
   loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)
   loss = tf.reduce_mean(loss)
   

3. 同步修改预测部分，因为现在是单输出，需要用sigmoid判断概率是否大于0.5来确定类别：

   predictions = {
       "classes": tf.cast(tf.greater(tf.nn.sigmoid(logits), 0.5), tf.int32),
       "probabilities": tf.nn.sigmoid(logits, name="sigmoid_tensor")
   }
   

3. 额外的避坑建议 #

• 输入数据归一化：检查你的图像数据是否已经缩放到0-1范围（比如除以255），未归一化的像素值（0-255）会导致模型权重更新异常，很容易出现NaN。
• 调试dropout：如果还是出现NaN，可以先把dropout的rate设为0，测试是否是dropout导致的训练不稳定。

内容的提问来源于stack exchange，提问作者Adrián Arroyo Perez