python label 与 one-hot 之间的互相转换

前言

有时候需要label，比如强化学习的离散动作空间，输出动作索引；有时候需要one-hot，比如训练数据或者输入上一个状态的动作，简单的互相转换还是重要的。

label 转 one-hot

通过 np.eye(action_dims)[actions] 快速生成：

>>> import numpy as np
>>> label = [1,2,2,3]
>>> np.eye(4)[label]
array([[0., 1., 0., 0.],
       [0., 0., 1., 0.],
       [0., 0., 1., 0.],
       [0., 0., 0., 1.]])

one-hot 转label

numpy可以通过 np.argmax(onehot, 1) 实现，pytorch 可以通过 torch.topk(one_hot, 1)[1].squeeze(1) 实现：

>>> import torch
>>> onehot
array([[0., 1., 0., 0.],
       [0., 0., 1., 0.],
       [0., 0., 1., 0.],
       [0., 0., 0., 1.]])
>>> np.argmax(onehot,1)
array([1, 2, 2, 3], dtype=int64)
>>> torch.topk(torch.tensor(onehot), 1)[1].squeeze(1)
tensor([1, 2, 2, 3])