TensorFlow孪生网络自定义对比损失函数梯度形状不匹配错误解决

问题背景 #

使用TensorFlow和Keras构建孪生网络，因数据量较大采用生成器实现批量加载数据，训练时使用自定义对比损失函数，梯度计算阶段出现形状不匹配错误。

错误信息 #

Epoch 1/10
y_true shape: (None, 1)
y_pred shape: (None, 1)
loss: Tensor("contrastive_loss/Mean:0", shape=(), dtype=float32)
y_true shape: (None, 1)
y_pred shape: (None, 1)
loss: Tensor("contrastive_loss/Mean:0", shape=(), dtype=float32)
   1/1125 [..............................] - ETA: 7:56:52 - loss: 0.1281 - accuracy: 0.0625
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-17-d373b6e0f1af> in <cell line: 11>()
      9 #val_gen = data_generator(val_pairs, val_labels, batch_size, os.path.join(extract_path, 'train'))
     10 
---> 11 model.fit(
     12     train_gen,
     13     # validation_data=val_gen,

1 frames
/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     51   try:
     52     ctx.ensure_initialized()
---> 53     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     54                                         inputs, attrs, num_outputs)
     55   except core._NotOkStatusException as e:

InvalidArgumentError: Graph execution error:

Detected at node 'gradient_tape/contrastive_loss/mul_1/BroadcastGradientArgs' defined at (most recent call last):
    File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
      return _run_code(code, main_globals, None,
------
------"Error messages condensed"
------
    File "/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/optimizer.py", line 276, in compute_gradients
      grads = tape.gradient(loss, var_list)
Node: 'gradient_tape/contrastive_loss/mul/BroadcastGradientArgs'
Incompatible shapes: [0,1] vs. [32,1]
     [[{{node gradient_tape/contrastive_loss/mul/BroadcastGradientArgs}}]] [Op:__inference_train_function_11308]

数据生成器代码 #

def data_generator(pairs, labels, batch_size, img_dir):
    """
    Generate batches of images and labels for training/validation.

    Parameters:
    - pairs: List of tuples containing left image id, list of candidate right image ids,
             and index of ground truth right image.
    - batch_size: Number of pairs to load in each batch.
    - img_dir: Directory containing the images.

    Yields:
    Batch of images and labels.
    """
    num_samples = len(pairs)

    while True:
        # Shuffle pairs for randomness in each epoch
        # np.random.shuffle(pairs)

        # Create a list of sequence indices
        indices = np.arange(num_samples)

        # Shuffle the indices
        np.random.shuffle(indices)

        # Use the shuffled indices to shuffle the sequences
        pairs = np.array(pairs)
        pairs = pairs[indices]
        pairs = pairs.tolist()

        labels = np.array(labels)
        labels = labels[indices]
        labels = labels.tolist()

        for start_idx in range(0, num_samples, batch_size):
            end_idx = min(start_idx + batch_size, num_samples)
            batch_pairs = pairs[start_idx:end_idx]
            labels = labels[start_idx:end_idx]

            left_images = []
            right_images = []
            # labels = []

            for pair in batch_pairs:
                left_img_id, right_img_id = pair

                # Load left image
                left_img = load_and_preprocess_image(left_img_id, img_dir, 'left')
                left_images.append(left_img)

                # Load right images
                right_img = load_and_preprocess_image(right_img_id, img_dir, 'right')
                right_images.append(right_img)

            # Convert lists to numpy arrays
            left_images = np.array(left_images)
            right_images = np.array(right_images)
            labels = np.array(labels)

            yield [left_images, right_images], labels

模型训练代码 #

batch_size = 32
train_gen = data_generator(train_pairs, train_labels, batch_size, os.path.join(extract_path, 'train'))
val_gen = data_generator(val_pairs, val_labels, batch_size, os.path.join(extract_path, 'train'))

model.compile(optimizer='rmsprop', loss=contrastive_loss)
model.fit(
    train_gen,
    validation_data=val_gen,
    steps_per_epoch=len(train_pairs) // batch_size,
    validation_steps=len(val_pairs) // batch_size,
    epochs=10
)

问题排查与解决 #

已验证前向传播时y_true和y_pred形状均为(None,1)，但报错出现[32,1]与[0,1]形状不匹配。经排查，问题出在数据生成器内部：生成器中错误覆盖了输入的labels变量，导致第二次迭代时无有效标签返回，引发形状不匹配。

修复方法：将生成器中的labels = labels[start_idx:end_idx]改为batch_labels = labels[start_idx:end_idx]，并将后续涉及该变量的代码同步替换为batch_labels，修改后的生成器关键部分如下：

for start_idx in range(0, num_samples, batch_size):
    end_idx = min(start_idx + batch_size, num_samples)
    batch_pairs = pairs[start_idx:end_idx]
    batch_labels = labels[start_idx:end_idx]  # 修改变量名避免覆盖

    left_images = []
    right_images = []

    for pair in batch_pairs:
        left_img_id, right_img_id = pair

        # Load left image
        left_img = load_and_preprocess_image(left_img_id, img_dir, 'left')
        left_images.append(left_img)

        # Load right images
        right_img = load_and_preprocess_image(right_img_id, img_dir, 'right')
        right_images.append(right_img)

    # Convert lists to numpy arrays
    left_images = np.array(left_images)
    right_images = np.array(right_images)
    batch_labels = np.array(batch_labels)  # 同步修改变量名

    yield [left_images, right_images], batch_labels  # 同步修改返回变量

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