如何使用 PyTorch 沿特定维度进行热编码? [英] How do I one hot encode along a specific dimension using PyTorch?

查看:42
本文介绍了如何使用 PyTorch 沿特定维度进行热编码?的处理方法,对大家解决问题具有一定的参考价值,需要的朋友们下面随着小编来一起学习吧!

问题描述

我有一个大小为 [3, 15, 136] 的张量,其中:

I have a tensor of size [3, 15, 136], where:

  • 3 是批量大小
  • 15 - 序列长度
  • 136 是令牌

我想使用 tokens 维度 (136) 中的概率对我的张量进行单热化.为此,我想提取序列长度中每个字母的标记维度并将 1 置于最大可能性并将所有其他标记标记为 0.

I want to one-hot my tensor using the probabilities in the tokens dimension (136). To do so I want to extract the tokens dimension for each letter in sequence length and put 1 to the largest possibility and mark all other tokens as 0.

推荐答案

你可以使用 PyTorch 的 one_hot 函数来实现:

You can use PyTorch's one_hot function to achieve this:

import torch.nn.functional as F

t = torch.rand(3, 15, 136)

F.one_hot(t.argmax(dim=2), 136)

这篇关于如何使用 PyTorch 沿特定维度进行热编码?的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!

查看全文
登录 关闭
扫码关注1秒登录
发送“验证码”获取 | 15天全站免登陆