FaceNet模型三元组生成报错问题及优化方案咨询

问题描述 #

我正在构建FaceNet人脸识别模型，需要生成三元组用于训练，但三元组生成代码出现报错，错误信息如下：

InvalidArgumentError: Shapes of all inputs must match: values[0].shape = [7,96,96,3] != values[2].shape = [6,96,96,3] [Op:Pack] name: stack

尝试使用ragged tensor但不知如何操作，希望得到代码修改建议或更优实现方案。

原三元组生成及测试代码 #

# 因为数据量很少，所以不使用原论文中的约束条件
def select_all_triplets(images, labels):
    batch_size = len(labels)
    pos_images = [] # 存储正样本图片
    neg_images = [] # 存储负样本图片
    
    for i in range(batch_size):
        anchor_label = labels[i]
        
        pos_list = [] # 存储正样本索引
        neg_list = [] # 存储负样本索引
        
        for j in range(batch_size):
            if j != i: # 因此 len(pos_list) + len(neg_list) = len(images) - 1
                if labels[j] == anchor_label:
                    pos_list.append(j)
                else:
                    neg_list.append(j)
                    
        pos_images.append(tf.gather(images, pos_list))
        neg_images.append(tf.gather(images, neg_list))
        print(pos_list, neg_list)
        
    positive_images = tf.random.shuffle(tf.stack(pos_images))
    negative_images = tf.random.shuffle(tf.stack(neg_images))
    
    return positive_images, negative_images

# 测试代码
img, lbl = next(iter(train_ds))
select_all_triplets(img, lbl)

报错原因 #

tf.stack要求所有输入张量的形状完全一致，但每个锚点对应的正/负样本数量不相同（比如部分锚点有7个正样本，部分只有6个），导致pos_images和neg_images中的张量形状不一致，无法完成堆叠操作。

解决方案 #

方案一：使用Ragged Tensor处理可变长度

Ragged Tensor支持存储形状不一致的张量，可直接替代普通张量完成堆叠操作：

import tensorflow as tf

def select_all_triplets_ragged(images, labels):
    batch_size = len(labels)
    pos_images = []
    neg_images = []
    
    for i in range(batch_size):
        anchor_label = labels[i]
        # 向量化筛选正/负样本索引，替代嵌套循环提升效率
        mask_pos = (labels == anchor_label) & (tf.range(batch_size) != i)
        mask_neg = (labels != anchor_label)
        
        pos_indices = tf.where(mask_pos)[:, 0]
        neg_indices = tf.where(mask_neg)[:, 0]
        
        pos_images.append(tf.gather(images, pos_indices))
        neg_images.append(tf.gather(images, neg_indices))
    
    # 转换为Ragged Tensor，支持可变长度张量堆叠
    positive_ragged = tf.ragged.stack(pos_images)
    negative_ragged = tf.ragged.stack(neg_images)
    
    # 对每个锚点对应的正/负样本单独打乱顺序
    positive_ragged = positive_ragged.map(tf.random.shuffle)
    negative_ragged = negative_ragged.map(tf.random.shuffle)
    
    return positive_ragged, negative_ragged

方案二：生成固定数量的三元组（更适合训练）

实际训练中不需要每个锚点对应所有正/负样本，通常每个锚点随机选1个正样本和1个负样本，生成形状固定的三元组，避免形状不一致问题：

def generate_fixed_triplets(images, labels):
    batch_size = tf.shape(labels)[0]
    triplets = []
    
    for i in range(batch_size):
        anchor_label = labels[i]
        # 获取当前锚点的所有正样本索引（排除自身）
        pos_indices = tf.where((labels == anchor_label) & (tf.range(batch_size) != i))[:, 0]
        # 获取所有负样本索引
        neg_indices = tf.where(labels != anchor_label)[:, 0]
        
        if tf.size(pos_indices) == 0 or tf.size(neg_indices) == 0:
            continue  # 跳过没有正/负样本的锚点
        
        # 随机选取一个正样本和一个负样本
        pos_idx = tf.random.shuffle(pos_indices)[0]
        neg_idx = tf.random.shuffle(neg_indices)[0]
        
        triplets.append([images[i], images[pos_idx], images[neg_idx]])
    
    # 返回形状固定的张量，可直接输入Triplet Loss计算
    return tf.stack(triplets)

方案优势说明 #

• 方案一适合需要保留所有正/负样本关联场景，但后续训练时需要额外处理Ragged Tensor的维度适配。
• 方案二生成的三元组形状固定（[num_triplets, 3, 96, 96, 3]），可直接对接FaceNet的Triplet Loss损失函数，更符合实际训练流程，是更推荐的实现方式。

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