TensorFlow中替代np.isin实现标签批量替换的无迭代方案

适配tf.function的纯TensorFlow实现方案 #

完全不需要Python侧迭代、也不需要调用numpy接口，所有逻辑都用TF原生算子实现，完美支持静态图加速：

import tensorflow as tf

# 原始输入定义
categories = [
 [0,2,4,6,8],
 [1,3,5,7,9]  
]
y = tf.constant([0, 1, 1, 2, 5, 4, 7, 9, 3, 3])

# 提前把分类规则转为TF常量，可放在@tf.function外部预定义
cat_tensor = tf.constant(categories, dtype=tf.int32)

@tf.function
def get_category_labels(y, cat_tensor):
    # 维度扩展适配广播机制：y的shape从[batch_size]转为[batch_size, 1, 1]
    y_expand = y[:, tf.newaxis, tf.newaxis]
    # 广播比较：每个y的元素和所有分类下的所有取值做相等判断
    # 输出shape为 [batch_size, 分类总数, 单个分类的取值数]
    equal_mask = tf.equal(y_expand, cat_tensor)
    # 对每个分类下的比较结果做或运算：只要该分类下有一个取值匹配，就标记为属于该分类
    # 输出shape为 [batch_size, 分类总数]
    cat_match = tf.reduce_any(equal_mask, axis=-1)
    # 取匹配成功的分类索引作为最终标签，转回和y一致的int32类型
    cat_labels = tf.cast(tf.argmax(cat_match, axis=-1), tf.int32)
    return cat_labels

# 验证输出
cat_labels = get_category_labels(y, cat_tensor)
print(cat_labels)
# 输出：tf.Tensor([0 1 1 0 1 0 1 1 1 1], shape=(10,), dtype=int32)

方案说明 #

1. 全程使用TF原生静态图算子，没有Python侧遍历张量、类型转换等动态操作，不会触发tf.function的迭代报错。
2. 利用张量广播机制替代嵌套循环，运算效率远高于逐元素遍历，batch越大性能优势越明显。
3. 如果不同分类的取值数量不一致，可以提前对categories做填充，填充值取不会出现在类标y中的值（比如-1）即可，不会影响匹配结果。
4. 如果存在y的取值不在所有分类中的场景，可以额外加判断逻辑：先通过tf.reduce_any(cat_match, axis=-1)找出无匹配的样本，做自定义异常处理即可。

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