将numpy从一次热转换为稀疏张量

Question

new_one_hot = None
for i in range(num_channels):
    new_channel = one_hot[:,:,i] * (i+1)
    if new_one_hot is None:
        new_one_hot = new_channel
    else:
        new_one_hot = np.dstack((new_one_hot, new_channel))

我正在尝试将上述代码转换为所有张量流操作，因为我正在使用它来计算损失函数。基本上，我正在尝试将单编码的标签转换为稀疏标签。

虽然我有点卡住了-如何将上述内容转换为张量流操作？我还能使用for循环吗？我看不到使用for循环的tensorflow很多。

如何建立索引和dstack？对此有些困惑。

这是我的以下尝试：

def convert_one_hot_to_sparse_tf(one_hot):
    num_channels = one_hot.get_shape()[2]
    new_one_hot = None
    for i in range(num_channels):
        new_channel = one_hot[:,:,i] * tf.constant(i+1)
        if new_one_hot is None:
            new_one_hot = new_channel
        else:
            new_one_hot = tf.stack([new_one_hot, new_channel], axis=-1)

    sparse_one_hot = tf.math.argmax(new_one_hot, axis=-1)
    return sparse_one_hot

但是我得到“形状必须是相同的等级，但必须是3和2 将形状0与其他形状合并。输入形状为[width，height，2]，[width，height]的'stack_3'（op：'Pack'）。

该函数的输入应为[width，height，num_channels] 我想要的输出将是[宽度，高度]