TensorFlow Keras训练验证码模型报错：张量与批次元素不匹配

问题描述 #

训练自定义验证码识别模型时持续出现如下报错：

Cannot add tensor to the batch: number of elements does not match. Shapes are: [tensor]: [13], [batch]: [5]
 [[{{node IteratorGetNext}}]] [Op:__inference_train_function_13370]

错误原因 #

核心问题是数据集内验证码标签长度不统一，执行batch操作时，TensorFlow尝试直接拼接不同长度的标签张量，导致形状不匹配。

解决方案 #

1. 统一标签长度（填充至最大长度） #

修改encode_single_sample函数，对编码后的标签进行填充，使其长度等于数据集的max_length：

def encode_single_sample(img_path, label):
    # 1. Read image
    img = tf.io.read_file(img_path)
    # 2. Decode and convert to grayscale
    img = tf.io.decode_png(img, channels=1)
    # 3. Convert to float32 in [0, 1] range
    img = tf.image.convert_image_dtype(img, tf.float32)
    # 4. Resize to the desired size
    img = tf.image.resize(img, [img_height, img_width], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
    # 5. Transpose the image because we want the time
    # dimension to correspond to the width of the image.
    img = tf.transpose(img, perm=[1, 0, 2])
    # 6. Map the characters in label to numbers
    label = char_to_num(tf.strings.unicode_split(label, input_encoding="UTF-8"))
    # 新增：将标签填充至最大长度，确保批量形状一致
    label = tf.pad(label, [[0, max_length - tf.shape(label)[0]]])
    # 7. Return a dict as our model is expecting two inputs
    return {"image": img, "label": label}

2. 修正CTCLayer的损失计算逻辑 #

原代码中标签长度计算错误，需修改为计算每个样本的真实标签长度（排除填充部分），同时修正输入序列长度的获取方式：

class CTCLayer(layers.Layer):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.loss_fn = keras.backend.ctc_batch_cost

    def call(self, y_true, y_pred):
        batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64")
        # 修正：RNN输出的序列长度是shape[1]
        input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64")
        # 计算每个样本的真实标签长度（统计非填充元素数量）
        label_length = tf.cast(tf.math.count_nonzero(y_true, axis=1), dtype="int64")

        input_length = input_length * tf.ones(shape=(batch_len, 1), dtype="int64")

        loss = self.loss_fn(y_true, y_pred, input_length, label_length)
        self.add_loss(loss)

        return y_pred

3. 调整模型标签输入的形状定义 #

在build_model函数中，将标签输入的形状改为固定的max_length，匹配填充后的标签形状：

def build_model():
    # Inputs to the model
    input_img = layers.Input(
        shape=(img_width, img_height, 1), name="image", dtype="float32"
    )
    # 修改标签输入为固定长度max_length
    labels = layers.Input(name="label", shape=(max_length,), dtype="float32")
    
    # 后续卷积、RNN等层代码保持不变
    # ...

原理说明 #

1. 标签填充：确保批量内所有样本的标签张量形状一致，解决批量拼接时的形状不匹配问题；
2. 真实长度计算：CTC损失需要知道每个样本的实际字符数量，而非填充后的长度，否则会导致损失计算错误；
3. 输入形状修正：让模型明确接收固定长度的标签张量，避免输入输出形状不兼容。

内容的提问来源于stack exchange，提问作者Iron Banker Of Braavos