TensorFlow：如何实现支持可变on_values且depth为张量的tf.one_hot

嘿，这个需求确实没法直接用原生tf.one_hot实现——它的depth参数必须是标量，而且on_value也只能设成统一的值，没法给每个索引单独指定不同的有效值。不过别担心，我们用几个基础TensorFlow操作组合一下就能搞定，思路其实很简单：先生成标准的one-hot掩码矩阵，再把掩码里的1替换成对应的depth值就行。

具体实现步骤 #

• 定义输入张量：把你的indices和depth定义为TensorFlow常量（或变量，完全支持动态张量）。这里把depth改名为depth_values，避免和tf.one_hot的参数名冲突。
• 确定输出类别数：也就是结果矩阵的列数，由indices中的最大值加1得到，确保能覆盖所有索引位置。
• 生成one-hot掩码：用原生tf.one_hot生成一个只有0和1的掩码矩阵，每个行对应indices的位置标记为1。
• 替换掩码中的1为对应depth值：把depth_values扩展成列向量，和掩码矩阵做元素相乘，这样掩码里的1就会被替换成对应的depth值，0保持不变。

完整代码示例 #

import tensorflow as tf

# 定义输入张量（支持动态张量，比如从其他计算节点传入）
indices = tf.constant([0, 1, 2], dtype=tf.int32)
depth_values = tf.constant([4, 1, 2], dtype=tf.float32)

# 计算输出的类别数（结果矩阵的列数）
num_classes = tf.reduce_max(indices) + 1

# 生成标准one-hot掩码矩阵
one_hot_mask = tf.one_hot(indices, depth=num_classes, dtype=tf.float32)

# 将depth值扩展为列向量，与掩码相乘得到最终结果
result = one_hot_mask * tf.expand_dims(depth_values, axis=1)

# 运行测试
# TensorFlow 2.x 直接打印结果
print(result.numpy())

# TensorFlow 1.x 需要使用Session运行
# with tf.Session() as sess:
#     print(sess.run(result))

输出结果 #

运行后会得到你期望的矩阵：

[[4. 0. 0.]
 [0. 1. 0.]
 [0. 0. 2.]]

这个方案完全支持depth_values作为动态张量，也能实现每个索引对应不同的on_value，完美匹配你的需求。

内容的提问来源于stack exchange，提问作者S. Salman