label轉onehot的很多,但是onehot轉label的有點難找,所以就只能自己實現以下,用的topk函數,不知道有沒有更好的實現
one_hot = torch.tensor([[0,0,1],[0,1,0],[0,1,0]]) print(one_hot) label = torch.topk(one_hot, 1)[1].squeeze(1) print(label)
tensor([[0, 0, 1],
[0, 1, 0],
[0, 1, 0]])
tensor([2, 1, 1])