如何使用PyTorch沿特定维度进行热编码？

Question

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

• 3 is batch size
• 15 - sequence length和
• 136 is tokens

我想用tokens维（136）中的概率来计算我的张量。为此，我想为序列长度中的每个字母提取tokens维度，并将1置于最大可能性，并将所有其他标记标记为0

Answer 1

您可以使用PyTorch的^{}函数来实现这一点：

import torch.nn.functional as F

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

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