如何在PyTorch中将整数tensor转换为对应索引位置为1的二进制tensor

你所需的是标准的独热编码转换操作，PyTorch内置了现成接口可以直接实现，不需要手动写循环逻辑。

实现示例 #

1. 输入构造 #

import torch
import torch.nn.functional as F

# 对应你的示例输入
input_tensor = torch.tensor([[1,3,2,6]])

2. 调用接口转换 #

# num_classes设置为输入最大值+1，刚好匹配你示例的输出维度需求
one_hot_tensor = F.one_hot(input_tensor, num_classes=input_tensor.max() + 1).squeeze(0)

3. 输出结果 #

运行后得到的one_hot_tensor和你期望的输出完全一致：

tensor([[0,1,0,0,0,0,0],
        [0,0,0,1,0,0,0],
        [0,0,1,0,0,0,0],
        [0,0,0,0,0,0,1]])

注意事项 #

• 如果你需要固定输出维度，不需要依赖输入的最大值，可以手动指定num_classes为你需要的数值即可。
• 输入张量必须为整数类型（int、long等），浮点类型张量调用接口会报错。
• 示例中添加squeeze(0)是因为你的输入为[1,4]的二维张量，one_hot转换后默认维度为[1,4,7]，去掉第0维即可得到你需要的[4,7]输出结构，你可以根据自己的输入维度按需调整维度处理逻辑。

内容的提问来源于stack exchange，提问作者Seraf Fej