One-hot function in PyTorch

最新推荐文章于 2026-07-02 20:53:08 发布

原创最新推荐文章于 2026-07-02 20:53:08 发布 · 1.7k 阅读

2 ·

本内容遵循CC 4.0 BY-SA版权协议

本文介绍并展示了四种在PyTorch中实现One-Hot编码的方法：使用scatter_函数、index_select()函数、Embedding模块以及自定义One_Hot模块。这些方法为不同场景下的One-Hot编码提供了灵活的选择。

Qwen3-32B-Chat 私有部署镜像 | RTX4090D 24G 显存 CUDA12.4 优化版

Qwen

文本生成

Qwen3

本镜像基于 RTX 4090D 24GB 显存 + CUDA 12.4 + 驱动 550.90.07 深度优化，内置完整运行环境与 Qwen3-32B 模型依赖，开箱即用。

One-hot function implementation in PyTorch

Method1: use scatter_ function

labels = [0, 1, 4, 7, 3, 2]
one_hot = torch.zeros(6, 8).scatter_(dim = 1, index = labels, value = 1)

Method2: use index_select() funtion

labels = [0, 1, 4, 7, 3, 2]
index = torch.eye(8)
one_hot = torch.index_select(index, dim = 0, index = labels)

Method3: use Embedding module

emb = nn.Embedding(8, 8)
emb.weight.data = torch.eye(8)

then we can get

emb(Variable(torch.LongTensor([1, 2], [3, 4])))
Variable containing:
(0,.,.) = 
0 1 0 0 0 0 0 0 
0 0 1 0 0 0 0 0
(1 ,.,.) =
0 0 0 1 0 0 0 0
0 0 0 0 1 0 0 0

Method4: create a module

class One_Hot(nn.Module):
    def __init__(self, depth):
        super(One_Hot,self).__init__()
        self.depth = depth
        self.ones = torch.sparse.torch.eye(depth)
    def forward(self, X_in):
        X_in = X_in.long()
        return Variable(self.ones.index_select(0,X_in.data))
    def __repr__(self):
        return self.__class__.__name__ + "({})".format(self.depth)