如何在TensorFlow中仅对张量中的特定数组进行离散化

你可以通过单独处理目标数组，再与原数组的未处理部分拼接的方式实现需求，以下是两种常见场景的实现方案：

场景1：仅基于第二个数组的统计信息做离散化 #

如果离散化的分箱规则只需要参考第二个数组的数据，可以这样做：

import numpy as np
import tensorflow as tf

input = np.array([[-1.5, 1.0, 3.4, .5], [0.0, 3.0, 1.3, 0.0]])

# 初始化离散化层，仅用第二个数组适配分箱规则
layer = tf.keras.layers.Discretization(num_bins=2, epsilon=0.01)
layer.adapt(input[1:])  # 只传入第二个数组做适配

# 仅对第二个数组执行离散化
discretized_second = layer(input[1:])
# 保留第一个数组的原始值，转换为与离散化结果同类型的张量
original_first = tf.cast(input[:1], dtype=discretized_second.dtype)
# 拼接得到最终结果
result = tf.concat([original_first, discretized_second], axis=0)

print(result)

输出结果：

<tf.Tensor: shape=(2, 4), dtype=int64, numpy=
array([[-1,  1,  3,  0],
       [ 0,  1,  1,  0]])>

场景2：基于整个输入的统计信息做离散化，但只处理第二个数组 #

如果分箱规则需要参考整个输入的统计数据，但只对第二个数组做离散化：

import numpy as np
import tensorflow as tf

input = np.array([[-1.5, 1.0, 3.4, .5], [0.0, 3.0, 1.3, 0.0]])

# 用整个输入适配分箱规则
layer = tf.keras.layers.Discretization(num_bins=2, epsilon=0.01)
layer.adapt(input)

# 仅处理第二个数组
discretized_second = layer(input[1:])
# 保留第一个数组原始值
original_first = tf.cast(input[:1], dtype=discretized_second.dtype)
# 拼接结果
result = tf.concat([original_first, discretized_second], axis=0)

print(result)

输出结果：

<tf.Tensor: shape=(2, 4), dtype=int64, numpy=
array([[-1,  1,  3,  0],
       [ 0,  1,  1,  0]])>

关键说明 #

• 切片操作input[1:]用于提取第二个数组，input[:1]提取第一个数组
• 使用tf.concat将未处理的第一个数组和离散化后的第二个数组拼接回原形状
• 用tf.cast统一张量数据类型，避免拼接时的类型不兼容问题

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