OneHot

基于给定的索引、深度和轴计算 one-hot 编码。在由 indices 输入定义的每个索引上设置 onValue，并在所有其他位置设置 offValue。

属性

axis：一个标量，指定添加输出 one-hot 编码的哪个维度。有效范围 [-rank(Indices)-1, rank(Indices)]

输入

indices：类型为 int32 的张量，包含热索引

values：类型为 T 的张量，包含 offValue 和 onValue

depth：类型为 int32 的张量，编码的深度

输出

output：类型为 T 的张量

数据类型

T：int8、int32、int64、float16、float32、bfloat16

形状信息

indices 是形状为 \([A_0,...,A_n]\) 的张量

values 是一个包含两个元素（rank=1）的张量，由 [offValue, onValue] 组成

depth 是一个 rank 为 0 的形状张量 \([d]\)

output 是形状为 \([A_0,...,A_{axis-1},d,A_{axis+1},...,A_n]\) 的张量

体积限制

indices 和 output 最多可以有 \(2^{31}-1\) 个元素。

示例

OneHot

inInd = network.add_input("indices", dtype=trt.int32, shape=(10,))
inVals = network.add_input("values", dtype=trt.float32, shape=(2,))
inDepth = network.add_input("depth", dtype=trt.int32, shape=())
axis = -1
depthVal = 4

opt_profile = get_runner.builder.create_optimization_profile()
opt_profile.set_shape_input("depth", [depthVal], [depthVal], [depthVal])
get_runner.config.add_optimization_profile(opt_profile)

layer = network.add_one_hot(inInd, inVals, inDepth, axis)
network.mark_output(layer.get_output(0))

inputs[inInd.name] = np.array([0, 1, 2, 3, 4, -1, -2, -3, -4, -5], dtype=np.int32)
inputs[inVals.name] = np.array([0.0, 3.14], dtype=np.float32)
inputs[inDepth.name] = np.array([depthVal], dtype=np.int32)

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array(
    [
        [3.14, 0.0, 0.0, 0.0],
        [0.0, 3.14, 0.0, 0.0],
        [0.0, 0.0, 3.14, 0.0],
        [0.0, 0.0, 0.0, 3.14],
        [0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 3.14],
        [0.0, 0.0, 3.14, 0.0],
        [0.0, 3.14, 0.0, 0.0],
        [3.14, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.0],
    ],
    dtype=np.float32,
)

C++ API

有关 C++ IOneHotLayer 算子的更多信息，请参阅 C++ IOneHotLayer 文档。

Python API

有关 Python IOneHotLayer 算子的更多信息，请参阅 Python IOneHotLayer 文档。