使用自定义TripletLoss训练时出现tf.function非首次调用创建变量错误

核心错误原因 #

1. 损失函数中硬编码了固定batch_size=64作为参数默认值，但训练过程中最后一个batch的样本数通常小于64，会导致生成的单位矩阵与相似度矩阵scores形状不匹配。
2. K.eye()在tf.function编译的计算图中接收固定Python整数作为尺寸参数时，一旦运行时遇到实际batch大小不等于64的情况，会触发计算图重编译，被识别为「非首次调用时创建变量」，触发对应的ValueError。

修复后的TripletLoss代码 #

def TripletLoss(y_true,y_pred, margin=0.25):
    # 动态获取当前batch的实际大小，兼容所有batch尺寸
    batch_size = tf.shape(y_pred)[0]
    v1, v2 = y_pred[:,:128],y_pred[:,-128:]
    scores = K.dot(v1, K.transpose(v2))
    positive = tf.linalg.diag_part(scores)
    # 生成对应尺寸的单位矩阵
    eye_mat = K.eye(batch_size)
    negative_without_positive = scores - 2 * eye_mat

    closest_negative = tf.reduce_max(negative_without_positive, axis=1)

    negative_zero_on_duplicate = scores * (1.0 - eye_mat)
    
    # 动态计算负样本均值，数值类型对齐避免报错
    mean_negative = K.sum(negative_zero_on_duplicate, axis=1) / tf.cast(batch_size-1, K.floatx())
    
    triplet_loss1 = K.maximum(0.0, margin - positive + closest_negative)
    
    triplet_loss2 = K.maximum(0.0, margin - positive + mean_negative)
    
    triplet_loss = K.mean(triplet_loss1 + triplet_loss2)

    return triplet_loss

补充说明 #

修复逻辑核心是将固定batch_size改为动态从输入张量的形状中获取，确保单位矩阵尺寸始终和当前batch的相似度矩阵匹配，避免计算图重复编译。如果你的数据集总量刚好是64的整数倍，原始代码可临时运行，但只要出现尺寸不一致的batch就会触发报错，修改后的代码可兼容所有场景。

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