Repository: pytorch/ignite

Issue:
Improve to_onehot to any shape
Idea is to improve `to_onehot` to accept any input shape: (N, ) and (N, H, W, D1, D2, ...).
Currently supported shape is `(N, )`. Thoughts ?

Discussion:
vfdev-5: @jasonkriss @anmolsjoshi please merge if OK

Patch:
diff --git a/ignite/utils.py b/ignite/utils.py
index d90af2bf..c9a9d3e7 100644
--- a/ignite/utils.py
+++ b/ignite/utils.py
@@ -42,6 +42,8 @@ def apply_to_type(input_, input_type, func):
 
 
 def to_onehot(indices, num_classes):
-    """Convert a tensor of indices to a tensor of one-hot indicators."""
-    onehot = torch.zeros(indices.size(0), num_classes, device=indices.device)
+    """Convert a tensor of indices of any shape `(N, ...)` to a
+    tensor of one-hot indicators of shape `(N, num_classes, ...)`.
+    """
+    onehot = torch.zeros(indices.shape[0], num_classes, *indices.shape[1:], device=indices.device)
     return onehot.scatter_(1, indices.unsqueeze(1), 1)


Test patch:
diff --git a/tests/ignite/test_utils.py b/tests/ignite/test_utils.py
index 5c119d4b..324489f1 100644
--- a/tests/ignite/test_utils.py
+++ b/tests/ignite/test_utils.py
@@ -39,3 +39,18 @@ def test_to_onehot():
     actual = to_onehot(indices, 4)
     expected = torch.eye(4)
     assert actual.equal(expected)
+
+    y = torch.randint(0, 21, size=(1000,))
+    y_ohe = to_onehot(y, num_classes=21)
+    y2 = torch.argmax(y_ohe, dim=1)
+    assert y.equal(y2)
+
+    y = torch.randint(0, 21, size=(4, 250, 255))
+    y_ohe = to_onehot(y, num_classes=21)
+    y2 = torch.argmax(y_ohe, dim=1)
+    assert y.equal(y2)
+
+    y = torch.randint(0, 21, size=(4, 150, 155, 4, 6))
+    y_ohe = to_onehot(y, num_classes=21)
+    y2 = torch.argmax(y_ohe, dim=1)
+    assert y.equal(y2)