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

【问题标题】：How do I one hot encode along a specific dimension using PyTorch?如何使用 PyTorch 沿特定维度进行热编码？
【发布时间】：2021-07-26 23:23:10
【问题描述】：

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

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

【问题讨论】：

【解决方案1】：

您可以使用 PyTorch 的 one_hot 函数来实现：

import torch.nn.functional as F

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

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

【讨论】：

谢谢！它似乎对我有用，但有一个例外。在此操作之前我掩盖了我的张量，因此填充字母在整个维度上为 0。不知何故，F.one_hot(torch.argmax(t, dim=2), 136) 选择了第一个 0 作为一个。我应该在“argmax”之后再次屏蔽它吗？
@julliet 是的，如果有多个最大值（例如全部为 0），则 argmax 返回第一个最大值（即第一个元素）的索引。
我可以在argmax 之后再次应用遮罩，对吗？
@julliet 是的 :)