One-hot encoding with label 0 all 0

Suppose I have a tensor a with size of (1, 1, 6, 6)

print(a)
tensor([[[[0., 0., 0., 0., 0., 0.],
          [0., 4., 4., 5., 5., 0.],
          [0., 4., 4., 5., 5., 0.],
          [0., 2., 2., 1., 1., 0.],
          [0., 2., 2., 1., 1., 0.],
          [0., 0., 0., 0., 0., 0.]]]])

and I have 5 classes in totall, is it possible to have one-hot encoding b with size of (1, 5, 6, 6) (without for loop if possible)

such that 0 for empty so one-hot for label 0 would be [0, 0, 0, 0, 0], and 1 would be [1, 0, 0, 0, 0]5 would be [0, 0, 0, 0, 1], I tried to use F.one_hot but cant get what I need. Thanks.

  • the label should be 0-4 since 5 classes, but I have 0 specifically for empty space so I change the label to 1-5.

You could zero out the first row:

x = torch.arange(5)
print(x)
# tensor([0, 1, 2, 3, 4])
y = F.one_hot(x, num_classes=5)
print(y)
# tensor([[1, 0, 0, 0, 0],
#         [0, 1, 0, 0, 0],
#         [0, 0, 1, 0, 0],
#         [0, 0, 0, 1, 0],
#         [0, 0, 0, 0, 1]])
y[:, 0] = 0
print(y)
# tensor([[0, 0, 0, 0, 0],
#         [0, 1, 0, 0, 0],
#         [0, 0, 1, 0, 0],
#         [0, 0, 0, 1, 0],
#         [0, 0, 0, 0, 1]])

so instead of using num_classes=5, I can use num_classes=6 and just discard the first channel by something like y[:, 1:]

This could also work assuming you want to use an indexing starting at 1 (which would make index 5 the 6th class in the default use case).