TensorFlow训练CNN时标签与Logits形状不匹配问题排查

首先咱们来拆解你遇到的报错：ValueError: Shape mismatch: The shape of labels (received (2,)) should equal the shape of logits except for the last dimension (received (1, 2)).

这个问题的核心有两点：

1. 损失函数与标签格式不匹配 #

你使用的tf.losses.sparse_softmax_cross_entropy是专门为整数类型的类别索引标签设计的（比如直接用0、1表示类别），但你提前对标签做了tf.one_hot编码，把它变成了二维的one-hot向量（形状[样本数, 类别数]），这直接导致了形状不匹配。

至于你提到的“传入模型前标签形状是(10,2)，模型里变成(2,)”：这是因为tf.data.Dataset.from_tensor_slices会把输入张量按第一个维度切片，每个样本的标签是长度为2的one-hot向量；而模型里通过tf.reshape给输入自动加上了batch维度，所以logits的形状是(1,2)，两者自然对不上。

2. 输入函数未设置Batch大小 #

你的输入函数没有对dataset设置batch，Estimator默认会逐个处理样本，这不仅会引发维度混乱，还会极大降低训练效率。

解决方案（两种可选，推荐第二种） #

方案一：保留one-hot标签，修改损失函数

如果你想继续用one-hot编码的标签，需要把损失函数换成适配one-hot格式的tf.losses.softmax_cross_entropy，同时调整评估指标的标签处理：

def cnn_model_fn(features, labels, mode):
    # ... 其他代码保持不变 ...
    
    # 替换损失函数为softmax_cross_entropy，接收one-hot标签
    loss = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logits)
    
    # ... 训练部分代码不变 ...
    
    # 评估指标需要把one-hot标签转回整数索引
    eval_metric_ops = {
        "accuracy": tf.metrics.accuracy(
            labels=tf.argmax(labels, axis=1), 
            predictions=predictions["classes"]
        )
    }
    return tf.estimator.EstimatorSpec(
        mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

同时记得给输入函数加上batch设置：

def my_input_fn():
    filenames = tf.constant(glob.glob("C:/test_proje/*.jpg"))
    labels = tf.constant([0, 0, 1, 1, 1, 1, 1, 0, 0, 0])
    labels = tf.one_hot(labels, 2)
    dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
    dataset = dataset.map(_parse_function)
    # 根据你的显存调整batch大小
    dataset = dataset.batch(batch_size=32)
    return dataset

方案二：去掉one-hot编码，使用整数标签（更高效）

sparse_softmax_cross_entropy比softmax_cross_entropy更节省内存、计算更高效，推荐直接使用原始整数标签：

1. 修改输入函数，移除tf.one_hot步骤并添加batch：

def my_input_fn():
    filenames = tf.constant(glob.glob("C:/test_proje/*.jpg"))
    # 保留原始整数标签，不做one-hot编码
    labels = tf.constant([0, 0, 1, 1, 1, 1, 1, 0, 0, 0])
    dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
    dataset = dataset.map(_parse_function)
    # 设置batch大小
    dataset = dataset.batch(batch_size=32)
    return dataset

2. 模型函数完全不需要修改！此时sparse_softmax_cross_entropy会完美匹配：labels的形状是[batch_size]（整数索引），logits的形状是[batch_size, 2]，完全符合函数要求，评估指标也能正常工作。

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