python - 如何将单热向量转换为多标签？

Question

我有一个多分类任务，并且我得到了一个热门类型的预测，比如

[[0, 1, 1],
 [0, 1, 0],
 [1, 0, 1]]

我希望将这个单热向量转换为标签

[[1, 2], [1], [0, 2]]

我已经尝试过 tf.argmax，但它不起作用。那么我该如何处理呢？

标签： pythonpython-3.xpython-2.7tensorflow

Answer 1

使用列表推导：

oheList = [[0, 1, 1],
           [0, 1, 0],
           [1, 0, 1]]
[[i for i in range(len(el)) if el[i]==1] for el in oheList]
# [[1, 2], [1], [0, 2]]