Pytorch将标签转为One-Hot编码

一、标签映射与One-Hot编码过程

先进行标签映射，要为每个分类建立一个整数索引，对于每个样本的标签，使用整数索引创建一个长度为类别总数的二进制向量。这个向量的所有元素都是0，除了与整数索引相对应的位置，该位置的值为1。

二、pytorch的官方实现

在pytorch中实现了one hot编码，就在torch.nn.functional里面，下面是它的注释当中的示例，我们开看看：

Examples:
    >>> F.one_hot(torch.arange(0, 5) % 3)
    tensor([[1, 0, 0],
            [0, 1, 0],
            [0, 0, 1],
            [1, 0, 0],
            [0, 1, 0]])
    >>> F.one_hot(torch.arange(0, 5) % 3, num_classes=5)
    tensor([[1, 0, 0, 0, 0],
            [0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0],
            [1, 0, 0, 0, 0],
            [0, 1, 0, 0, 0]])
    >>> F.one_hot(torch.arange(0, 6).view(3,2) % 3)
    tensor([[[1, 0, 0],
             [0, 1, 0]],
            [[0, 0, 1],
             [1, 0, 0]],
            [[0, 1, 0],
             [0, 0, 1]]])

我们可以根据那自己实现的与它给出的这个示例进行比对，一样就当然没问题了。

三、手写实现

首先，在原先的函数（one_hot）当中numclass=-1，类别当然不能为1，说明这里是自动进行了计算，大家普遍使用的方式都是创建一个全零矩阵，使用 scatter_ 函数进行独热编码，作用是按照给定的索引，在指定的维度上进行赋值。

def one_hot(labels, num_classes=-1):
    """
    将标签转为独热编码, 经过测试与torch.nn.functional里面的函数测试相同
    :param labels: 标签
    :param num_classes: 默认为-1, 表示进行自动计算类别最大的那个
    Examples:
        >>> label_1 = torch.arange(0, 5) % 3
        # tensor([0, 1, 2, 0, 1])
        >>> label_2 = torch.arange(0, 6).view(3, 2) % 3
        # tensor([[0, 1], [2, 0], [1, 2]])
        >>> print(one_hot(label_1))
        tensor([[1, 0, 0],
                [0, 1, 0],
                [0, 0, 1],
                [1, 0, 0],
                [0, 1, 0]])
        >>> print(one_hot(label_1, 5))
        tensor([[1, 0, 0, 0, 0],
                [0, 1, 0, 0, 0],
                [0, 0, 1, 0, 0],
                [1, 0, 0, 0, 0],
                [0, 1, 0, 0, 0]])
        >>> print(one_hot(label_2))
        tensor([[[1, 0, 0],
                 [0, 1, 0]],
                [[0, 0, 1],
                 [1, 0, 0]],
                [[0, 1, 0],
                 [0, 0, 1]]])
    """
    if num_classes == -1:
        num_classes = int(labels.max()) + 1
    one_hot_tensor = torch.zeros(labels.size() + (num_classes,), dtype=torch.int64)
    one_hot_tensor.scatter_(-1, labels.unsqueeze(-1).to(torch.int64), 1)
    return one_hot_tensor


label_1 = torch.arange(0, 5) % 3
# tensor([0, 1, 2, 0, 1])
label_2 = torch.arange(0, 6).view(3, 2) % 3
# tensor([[0, 1], [2, 0], [1, 2]])
print(one_hot(label_1))
print(one_hot(label_1, 5))
print(one_hot(label_2))

首先是判断分类数是不是为-1，如果是就根据其中的最大值+1进行自动计算。然后创建一个契合分类数量的全零矩阵。

在这里，labels.unsqueeze(-1)用于在标签的最后一个维度上添加一个维度，以便与独热编码张量进行广播操作。

假设原始的 labels 张量的形状为 (batch_size,)，那么经过 unsqueeze(-1) 操作后，形状变为 (batch_size, 1)。这样，每个样本的标签都被表示为一个列向量，而不再是一个标量。scatter_函数在最后一个维度进行操作，也就是对类别总数的维度进行操作，而 1 是要赋给相应位置的值。

labels.unsqueeze(-1) 已经确保了与 one_hot_tensor 的形状匹配，所以在这里能够正确地进行广播和赋值操作。

下面这一种是应用于分割网络当中，在保留输入标签张量形状的同时，将独热编码张量的最后一个维度设置为分类数num_classes，确保独热编码张量与输入标签张量具有相同的形状。

def get_one_hot(labels, num_classes=-1):
    """用于分割网络的one hot"""
    labels = torch.as_tensor(labels)
    ones = one_hot(labels, num_classes)
    return ones.view(*labels.size(), num_classes)


if __name__=="__main__":
    seg_labels = torch.randint(0, 3, size=[512, 512])
    print(get_one_hot(seg_labels))
    print(get_one_hot(seg_labels).shape)   # torch.Size([512, 512, 3])

你可以将这里应用于自定义dataset部分。