torch.nn.functional.one_hot should gracefully skip negative and out-of-range indices

[msbaines]

🚀 Feature

There are useful algorithms that benefit for one_hot skipping negative and out-of-range indices. tensorflow does not crash in these scenarios and instead emits a 0 vector.

>>> x = torch.tensor([1, 2, 8])
>>> F.one_hot(x, num_classes=4)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Class values must be smaller than num_classes.

>>> x = torch.tensor([1, 2, -1])
>>> F.one_hot(x, num_classes=4)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Class values must be non-negative.

Motivation

Be able to simply represent some algorithms.

Pitch

Be able to simply represent some algorithms. Easy to copy code/algorithms to/from tf and torch.

Alternatives

Explicitly clear negative and out-of-range values before calling one_hot and then cleanup the one_hot vector after. This is less efficient and less elegant.

Additional context

cc @albanD @mruberry

[soumith]

the errors should be the default, but let's say F.one_hot(..., relaxed=True) or something can do the more relaxed behavior

[albanD]

Also having negative indices being nothing and not the last element looks very un-intuitive to me given the rest of the code base (and python in general)!

[msbaines]

Also having negative indices being nothing and not the last element looks very un-intuitive to me given the rest of the code base (and python in general)!

The tensorflow behavior is to set to zero so I would recommend that for consistency. (It actuals feels more intuitive to me to zero but I see your point about pythonic -1 indexing)

https://colab.research.google.com/drive/1CXtpfGLxetqoQIP1OxfIIxt9EIVP7MpF?usp=sharing

[ssnl]

maybe let's just have options like out_of_bounds='strict'/'relaxed'/'wrapping'

[ngimel]

Should we have same options for indexing? For scatter/gather? For embedding?

[gchanan]

if we were going to follow PyTorch patterns, wouldn't having an ignore_index (or similar -- zero_hot_index?) a-la NLLLoss be more consistent? Or just zero out the entries you don't want (we can demonstrate how to do that in the documentation/error message)

[mruberry]

I recommend wrapping the PyTorch function to get the functionality you want.

I'm leery of cluttering the UX with lots of flags, and I worry that users might be confused the other way ("I accidentally put in a negative index and it didn't error!") if we change the default behavior.

[visualDust]

Hi, I've run into this issue and currently I'm getting some workaround. I strongly agree with @mruberry and I solved this problem by warping a more complex one hot encoding. Hope my code helps.

Code

Code below drops all negative values

def one_hot_encode(
    mask: torch.Tensor,
    num_classes: int,
    ignored_label: Union[str, int] = "negative",
):
    """Convert the mask to a one-hot encoded representation by @visualDust

    Args:
        mask (torch.Tensor): indexed label image. Should types int
        num_classes (int): number of classes
        ignored_label (Union[str|int], optional): specify labels to ignore, or ignore by pattern. Defaults to "negative".

    Returns:
        torch.Tensor: one hot encoded tensor
    """
    original_shape = mask.shape
    for _ in range(4 - len(mask.shape)):
        mask = mask.unsqueeze(0)  # H W -> C H W -> B C H W, if applicable
    # start to handle ignored label
    # convert ignored label into positive index bigger than num_classes
    if type(ignored_label) is int:
        mask[mask == ignored_label] = num_classes
    elif ignored_label == "negative":
        print(mask.shape, mask.min())
        mask[mask < 0] = num_classes
        print(mask.shape, mask.min())

    # check if mask image is valid
    if torch.max(mask) > num_classes:
        raise RuntimeError("class values must be smaller than num_classes.")
    B, _, H, W = mask.shape
    one_hot = torch.zeros(B, num_classes + 1, H, W)
    one_hot.scatter_(1, mask, 1)  # mark 1 on channel(dim=1) with index of mask
    one_hot = one_hot[:, :num_classes] # remove ignored label(s)
    for _ in range(len(one_hot.shape) - len(original_shape)):
        one_hot.squeeze_(0)  # B C H W -> H W ->  C H W, if applicable
    return one_hot

Example usage

Purpose: convert a cityscapes label image into per class image

... read label image and map all ignored label to -1 # do label mapping
plt.imshow(label)
gt_tensor = torch.tensor(label).type(torch.int64)

what the label image looks like:

converting into per class tensor:

plt.figure(figsize=(12,10))
sub_idx = 1
for b in one_hot_encode(gt_tensor, num_classes=19):
    plt.subplot(6, 4, sub_idx)
    plt.imshow(b)
    plt.axis('off')
    sub_idx += 1
plt.show()

returns a 19-channel tensor after convert: