pytorch将标签转为onehot

由于想多分类中使用Diceloss，所以需要将[0,1,2,..N]类型的标签转化为onehot类型。

1、在cpu上处理

input数据类型： torch.LongTensor()

        数据形状：[bs, 1, *]         可为2D或3D数据      

def make_one_hot(input, num_classes):
    """Convert class index tensor to one hot encoding tensor.
    Args:
         input: A tensor of shape [bs, 1, *]
         num_classes: An int of number of class
    Returns:
        A tensor of shape [bs, num_classes, *]
    """
    shape = np.array(input.shape)
    shape[1] = num_classes
    shape = tuple(shape)
    result = torch.zeros(shape)
    result = result.scatter_(1, input.cpu(), 1)

    return result

 

2、在GPU上处理

input数据类型： torch.LongTensor().cuda()

        数据形状：[bs, 1, *]         可为2D或3D数据 

def make_one_hot(input, num_classes):
    """Convert class index tensor to one hot encoding tensor.
    Args:
         input: A tensor of shape [bs, 1, *]
         num_classes: An int of number of class
    Returns:
        A tensor of shape [bs, num_classes, *]
    """
    shape = np.array(input.shape)
    shape[1] = num_classes
    shape = tuple(shape)
    result = torch.zeros(shape).cuda()
    result = result.scatter_(1, input, 1)

    return result

3、温馨提示

1、FloatTensor转化为LongTensor：

# 此时的输入label为FloatTensor，可在cuda，也可是cpu
label_long = label.long()

2、 Tensor增加一个维度

label_onehot = label_onehot.unsqueeze(1)   #在第一维增加一个维度

3、多分类交叉熵是不需要将标签转为onehot的

详情请查看  https://blog.csdn.net/longshaonihaoa/article/details/105253553

4、最近版pytorch有直接的转化为onehot的代码，了解之后更新。