如何在TensorFlow中高效查找两个张量的相交行？

针对大规模二维张量，直接用tf.equal结合广播会生成(len(tensor1), len(tensor2))规模的中间矩阵，内存占用爆炸。下面提供两种内存与时间双优化的实现方案：

方案一：基于哈希表的快速查询（推荐） #

核心思路是将tensor2的每行映射为唯一可哈希的键（比如字符串），存入哈希表后，直接查询tensor1的每行是否在哈希表中。时间复杂度O(n+m)，内存复杂度O(n+m)，完全避免广播带来的内存开销。

代码实现 #

import tensorflow as tf

def rows_exist(tensor1, tensor2):
    # 将张量的每行转换为唯一字符串键
    def tensor_to_row_strings(tensor):
        # 把每行元素转为字符串后拼接，确保不同行生成不同字符串
        return tf.strings.reduce_join(
            tf.as_string(tensor, precision=10),  # 浮点数需指定精度避免匹配错误
            separator=',',
            axis=1
        )
    
    # 构建tensor2的行哈希表，默认返回False（表示不存在）
    tensor2_rows = tensor_to_row_strings(tensor2)
    hash_table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(tensor2_rows, tf.ones_like(tensor2_rows, dtype=tf.bool)),
        default_value=False
    )
    
    # 查询tensor1的每行是否存在
    tensor1_rows = tensor_to_row_strings(tensor1)
    return hash_table.lookup(tensor1_rows)

# 测试示例
tensor1 = tf.constant([[0,1,1], [0,1,0], [0,1,2]])
tensor2 = tf.constant([[0,0,0],[0,0,1],[0,1,1],[1,1,1]])
print(rows_exist(tensor1, tensor2).numpy())  # 输出: [ True False False]

注意事项 #

• 若张量元素为浮点数，必须通过precision参数指定足够的精度，避免因浮点精度丢失导致的错误匹配。
• 哈希表初始化仅需执行一次，若tensor2固定不变，可提前初始化复用，进一步提升效率。

方案二：基于排序与二分搜索的实现 #

如果不想用哈希表，也可以通过将张量行扁平化、排序后结合二分搜索实现：

代码实现 #

import tensorflow as tf

def rows_exist_via_sort(tensor1, tensor2):
    # 扁平化每行，将二维行转为一维标量（仅适用于整数张量，浮点数需谨慎）
    def flatten_rows(tensor):
        # 计算每行的"唯一标识值"，比如用基数转换的思路
        rank = tf.shape(tensor)[1]
        bases = tf.pow(10, tf.range(rank-1, -1, -1), dtype=tensor.dtype)
        return tf.tensordot(tensor, bases, axes=1)
    
    # 处理tensor2并排序
    tensor2_flat = flatten_rows(tensor2)
    sorted_tensor2 = tf.sort(tensor2_flat)
    
    # 处理tensor1并二分搜索
    tensor1_flat = flatten_rows(tensor1)
    # 搜索每个元素是否存在于排序后的tensor2中
    _, idx = tf.searchsorted(sorted_tensor2, tensor1_flat, side='left')
    idx = tf.clip_by_value(idx, 0, tf.shape(sorted_tensor2)[0]-1)
    exists = tf.equal(sorted_tensor2[idx], tensor1_flat)
    return exists

# 测试整数张量示例
tensor1 = tf.constant([[0,1,1], [0,1,0], [0,1,2]])
tensor2 = tf.constant([[0,0,0],[0,0,1],[0,1,1],[1,1,1]])
print(rows_exist_via_sort(tensor1, tensor2).numpy())  # 输出: [ True False False]

局限性 #

• 仅适合整数张量，浮点数容易因精度问题生成重复标识，导致匹配错误。
• 基数的选择需确保不会溢出，若张量元素范围大，需改用更大的基数或其他扁平化方式。

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