how to convert one-hot vectors to multi-label?

Viewed 181

I have a multi-classification task, and I have gotten the one-hot type predictions like

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

I wish to convert this one-hot vector to labels like

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

I have tried tf.argmax, but it doesn't work. So how can I deal with it?

2 Answers

Using list comprehensions:

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]]

Another way to solve the problem,

import numpy as np

arr = np.array([[0, 1, 1],
                [0, 1, 0],
                [1, 0, 1]])

result = {}
for r, c in zip(*np.where(arr == 1)):
    result.setdefault(r, []).append(c)

print(result.values())
[[1, 2], [1], [0, 2]]
Related