如何在TensorFlow中创建指定深度、指定索引处为1的掩码张量？

嘿，我来帮你搞定这个张量生成的需求！你要做的是根据给定的索引张量，生成一个指定深度的张量，把每个样本对应的索引位置设为1，其余为0对吧？先看你的示例：

示例说明 #

输入的索引张量：

[[1 3] [2 4] [0 4]]

指定深度depth=5，输出的目标张量：

[[0. 1. 0. 1. 0.]
 [0. 0. 1. 0. 1.]
 [1. 0. 0. 0. 1.]]

这本质上是多标签的one-hot编码场景，下面给你两种主流框架的实现方案，都是高效且易理解的：

实现方案（PyTorch） #

方法1：直观的索引赋值 #

如果你的张量规模不大，直接循环赋值非常直观：

import torch

# 输入索引张量
indices = torch.tensor([[1, 3], [2, 4], [0, 4]])
depth = 5

# 创建形状为 (样本数, depth) 的全零张量
output = torch.zeros(indices.shape[0], depth, dtype=torch.float32)

# 遍历每个样本，把对应索引位置设为1
for sample_idx, idx_list in enumerate(indices):
    output[sample_idx, idx_list] = 1.0

print(output)

方法2：高效的scatter_方法 #

当处理大规模张量时，用PyTorch内置的scatter_方法会比循环快很多，它专门用来按索引填充值：

import torch

indices = torch.tensor([[1, 3], [2, 4], [0, 4]])
depth = 5

output = torch.zeros(indices.shape[0], depth, dtype=torch.float32)
# dim=1表示按列维度填充，把indices指定的位置设为1.0
output.scatter_(1, indices, 1.0)

print(output)

实现方案（TensorFlow） #

如果你用TensorFlow，同样可以用张量散射更新的方式实现：

import tensorflow as tf

indices = tf.constant([[1, 3], [2, 4], [0, 4]])
depth = 5

# 创建全零张量
output = tf.zeros((tf.shape(indices)[0], depth), dtype=tf.float32)
# 构造散射更新的索引：每个位置是 (样本下标, 索引值)
scatter_positions = tf.concat([tf.expand_dims(tf.range(tf.shape(indices)[0]), 1), indices], axis=1)
# 更新指定位置为1
output = tf.tensor_scatter_nd_update(output, scatter_positions, tf.ones(tf.size(indices), dtype=tf.float32))

print(output.numpy())

核心思路其实很简单：先创建一个符合要求形状的全零张量，然后精准定位需要设为1的位置，把这些位置的值替换掉就行。用框架内置的散射方法能避免Python循环的开销，适合处理大张量。

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