Индексный массив для горячей матрицы идентификаторов

Question

Предположим, у меня есть массив [0, 2], я хотел бы вывести матрицу, состоящую из одного горячего вектора на основе [0, 2], например [[1, 0, 0] [0, 0, 1] ] (Обратите внимание, что второе измерение выходной матрицы предполагается равным 3, но оно может быть любым числом, превышающим argmax ([0,2]), равным 2.

Я могу думать только о том, как этого добиться функция. Есть ли более простой способ.

t = torch.tensor([0,2])
dim2_size = 3
id_t = torch.zeros(t.shape[0], dim2_size)
row_idx = 0
for i in t:
  col_idx = i.item()
  id_t[row_idx, col_idx] = 1
  row_idx += 1
id_t

Answer 1

метод через numpy проще

import torch
import numpy as np

labels =[0,2]
output=np.eye(max(labels)+1)[labels]
print(torch.from_numpy(output))

Answer 2

В Pytorch это лучше всего сделать с помощью scatter_.

t = torch.tensor([0,2]).unsqueeze(0)
num_dims = 3
id_t = torch.zeros(num_dims, t.shape[1]).scatter_(0, t, 1)

Это дает вам id_t как:

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

Answer 3

Этот не использует l oop.

import torch

labels = torch.tensor([0, 2])
one_hot = torch.zeros(labels.shape[0], torch.max(labels)+1)
one_hot[torch.arange(labels.shape[0]), labels] = 1

print(one_hot)

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