Convert int into one-hot format

Run into the issue myself and did some searching, torch.sparse.torch.eye(num_labels).index_select(dim=0, index=labels) also seems to work pretty well in addition to the scatter_ solution in the 0.3 release.


def get_one_hot(preds,gt):
encoded_target =
target = gt.unsqueeze(1) # now target is in shape [BCHW]=[20,1,240,240]
unseq = target.long()
unseq =

# encoded_target.scatter_(dim,index,val)
# unseq dim 'dim' must be 1
encoded_target.scatter_(1, unseq, 1)
return encoded_target

It returns the one hot encoding of the target. In my case the target was of shape [1,1,240,240] and preds of shape [1,5,240,240]

Here’s a tensorflow-like solution based on previous code in this thread

def to_one_hot(y, n_dims=None):
    """ Take integer y (tensor or variable) with n dims and convert it to 1-hot representation with n+1 dims. """
    y_tensor = if isinstance(y, Variable) else y
    y_tensor = y_tensor.type(torch.LongTensor).view(-1, 1)
    n_dims = n_dims if n_dims is not None else int(torch.max(y_tensor)) + 1
    y_one_hot = torch.zeros(y_tensor.size()[0], n_dims).scatter_(1, y_tensor, 1)
    y_one_hot = y_one_hot.view(*y.shape, -1)
    return Variable(y_one_hot) if isinstance(y, Variable) else y_one_hot


It is also possible to abuse broadcasting and do:

# some labels
labels = torch.arange(3)
labels = labels.reshape(3, 1)

num_classes = 4
one_hot_target = (labels == torch.arange(num_classes).reshape(1, num_classes)).float()


 1  0  0  0
 0  1  0  0
 0  0  1  0
[torch.FloatTensor of size (3,4)]

You can use torch.eye function for this:

def one_hot_embedding(labels, num_classes):
    """Embedding labels to one-hot form.

      labels: (LongTensor) class labels, sized [N,].
      num_classes: (int) number of classes.

      (tensor) encoded labels, sized [N, #classes].
    y = torch.eye(num_classes) 
    return y[labels] 

This should help!


As the reply by Adam Paszke on May 17 suggested, this is not a good method since you might need to create a large weight matrix everytime and store it in memory.

Wow. This is way simpler than what I used before. I like this one. Thanks!

how can i convert integer values i.e. 0,1,2,3 into the form of one hot vector i.e. [1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]?

@apaszke do you think an api e.g. torch.create_one_hot() will be useful and worth providing?
It will take input 1-d tensor and create a 2-d tensor with respective to input encoding

This code should work:

idx = torch.tensor([0, 1, 2, 3])
torch.zeros(len(idx), idx.max()+1).scatter_(1, idx.unsqueeze(1), 1.)

Thank you, It worked :slight_smile:

1 Like

Finally torch.nn.functional.one_hot has been added.

>>> import torch
>>> x = torch.arange(0,5)
>>> x
tensor([0, 1, 2, 3, 4])
>>> torch.nn.functional.one_hot(x)
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]])


A batch version that works:
batch_size = 3
y = torch.randint(0, 3, (batch_size, 5)).type(torch.LongTensor)
y_onehot = torch.FloatTensor(batch_size, 5, 5)
y_onehot.scatter_(2, torch.unsqueeze(y, 2), 1)

Thanks, how do you cope with ignore label please ?

I would say don’t bother. The only ways you’ll ever use those one-hot variables is either to embed them (in which case nn.Embedding allows you to do so directly from the indices) or use them in a loss function, in which case why not use a loss function that takes the indices directly.

Are you sure about this? I am currently calling it like so:

loss(predictions, target_int_index) and it is complaining.

(predictions is 1xK and target_int_index is a scalar)

Sorry to revive an old topic. But this is the first result I get on google when searching for one hot encoding in PyTorch. I just want to mention that one can use torch.nn.functional.one_hot ( for this.


@cossio torch.nn.functional.one_hot doesn’t allow specifying a dimension though…

F.one_hot will accept a varying number of dimensions and will append an additional dimension in the shape nb_classes.
You could permute this tensor afterwards, if necessary.

@ptrblck You are right, one can transpose the tensor, apply F.one_hot, then transpose back, but then the solution with scatter_ seems to be simpler.

1 Like