Repository: Project-MONAI/MONAI

Issue:
Dimension is hardcoded for softmax/argmax for PostTransform
For both Dictionary and Array based post transform, the dimension is hardcoded.

Softmax:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/transforms/post/array.py#L115

Argmax:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/transforms/post/array.py#L203

When the input is `[Batch, Channel, x, y...]` these transforms can't be used as dimension is kind of hardcoded here..

> Possibly similar things can be fixed in other places in POST Transform.. (when you search `dim=0`)



Discussion:

Patch:
diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py
--- a/monai/transforms/post/array.py
+++ b/monai/transforms/post/array.py
@@ -55,7 +55,7 @@
 
 class Activations(Transform):
     """
-    Add activation operations to the model output, typically `Sigmoid` or `Softmax`.
+    Activation operations, typically `Sigmoid` or `Softmax`.
 
     Args:
         sigmoid: whether to execute sigmoid function on model output before transform.
@@ -64,6 +64,8 @@ class Activations(Transform):
             Defaults to ``False``.
         other: callable function to execute other activation layers, for example:
             `other = lambda x: torch.tanh(x)`. Defaults to ``None``.
+        kwargs: additional parameters to `torch.softmax` (used when ``softmax=True``).
+            Defaults to ``dim=0``, unrecognized parameters will be ignored.
 
     Raises:
         TypeError: When ``other`` is not an ``Optional[Callable]``.
@@ -72,9 +74,12 @@ class Activations(Transform):
 
     backend = [TransformBackends.TORCH]
 
-    def __init__(self, sigmoid: bool = False, softmax: bool = False, other: Optional[Callable] = None) -> None:
+    def __init__(
+        self, sigmoid: bool = False, softmax: bool = False, other: Optional[Callable] = None, **kwargs
+    ) -> None:
         self.sigmoid = sigmoid
         self.softmax = softmax
+        self.kwargs = kwargs
         if other is not None and not callable(other):
             raise TypeError(f"other must be None or callable but is {type(other).__name__}.")
         self.other = other
@@ -112,7 +117,7 @@ def __call__(
         if sigmoid or self.sigmoid:
             img_t = torch.sigmoid(img_t)
         if softmax or self.softmax:
-            img_t = torch.softmax(img_t, dim=0)
+            img_t = torch.softmax(img_t, dim=self.kwargs.get("dim", 0))
 
         act_func = self.other if other is None else other
         if act_func is not None:
@@ -123,12 +128,11 @@ def __call__(
 
 class AsDiscrete(Transform):
     """
-    Execute after model forward to transform model output to discrete values.
-    It can complete below operations:
+    Convert the input tensor/array into discrete values, possible operations are:
 
-        -  execute `argmax` for input logits values.
-        -  threshold input value to 0.0 or 1.0.
-        -  convert input value to One-Hot format.
+        -  `argmax`.
+        -  threshold input value to binary values.
+        -  convert input value to One-Hot format (set ``to_one_hot=N``, `N` is the number of classes).
         -  round the value to the closest integer.
 
     Args:
@@ -140,6 +144,9 @@ class AsDiscrete(Transform):
             Defaults to ``None``.
         rounding: if not None, round the data according to the specified option,
             available options: ["torchrounding"].
+        kwargs: additional parameters to `torch.argmax`, `monai.networks.one_hot`.
+            currently ``dim``, ``keepdim``, ``dtype`` are supported, unrecognized parameters will be ignored.
+            These default to ``0``, ``True``, ``torch.float`` respectively.
 
     Example:
 
@@ -165,6 +172,7 @@ def __init__(
         to_onehot: Optional[int] = None,
         threshold: Optional[float] = None,
         rounding: Optional[str] = None,
+        **kwargs,
     ) -> None:
         self.argmax = argmax
         if isinstance(to_onehot, bool):  # for backward compatibility
@@ -172,6 +180,7 @@ def __init__(
         self.to_onehot = to_onehot
         self.threshold = threshold
         self.rounding = rounding
+        self.kwargs = kwargs
 
     def __call__(
         self,
@@ -200,13 +209,15 @@ def __call__(
         img = convert_to_tensor(img, track_meta=get_track_meta())
         img_t, *_ = convert_data_type(img, torch.Tensor)
         if argmax or self.argmax:
-            img_t = torch.argmax(img_t, dim=0, keepdim=True)
+            img_t = torch.argmax(img_t, dim=self.kwargs.get("dim", 0), keepdim=self.kwargs.get("keepdim", True))
 
         to_onehot = self.to_onehot if to_onehot is None else to_onehot
         if to_onehot is not None:
             if not isinstance(to_onehot, int):
-                raise AssertionError("the number of classes for One-Hot must be an integer.")
-            img_t = one_hot(img_t, num_classes=to_onehot, dim=0)
+                raise ValueError("the number of classes for One-Hot must be an integer.")
+            img_t = one_hot(
+                img_t, num_classes=to_onehot, dim=self.kwargs.get("dim", 0), dtype=self.kwargs.get("dtype", torch.float)
+            )
 
         threshold = self.threshold if threshold is None else threshold
         if threshold is not None:
@@ -217,7 +228,7 @@ def __call__(
             look_up_option(rounding, ["torchrounding"])
             img_t = torch.round(img_t)
 
-        img, *_ = convert_to_dst_type(img_t, img, dtype=torch.float)
+        img, *_ = convert_to_dst_type(img_t, img, dtype=self.kwargs.get("dtype", torch.float))
         return img
 
 
diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py
--- a/monai/transforms/post/dictionary.py
+++ b/monai/transforms/post/dictionary.py
@@ -104,6 +104,7 @@ def __init__(
         softmax: Union[Sequence[bool], bool] = False,
         other: Optional[Union[Sequence[Callable], Callable]] = None,
         allow_missing_keys: bool = False,
+        **kwargs,
     ) -> None:
         """
         Args:
@@ -117,6 +118,8 @@ def __init__(
                 for example: `other = torch.tanh`. it also can be a sequence of Callable, each
                 element corresponds to a key in ``keys``.
             allow_missing_keys: don't raise exception if key is missing.
+            kwargs: additional parameters to `torch.softmax` (used when ``softmax=True``).
+                Defaults to ``dim=0``, unrecognized parameters will be ignored.
 
         """
         super().__init__(keys, allow_missing_keys)
@@ -124,6 +127,7 @@ def __init__(
         self.softmax = ensure_tuple_rep(softmax, len(self.keys))
         self.other = ensure_tuple_rep(other, len(self.keys))
         self.converter = Activations()
+        self.converter.kwargs = kwargs
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
@@ -147,6 +151,7 @@ def __init__(
         threshold: Union[Sequence[Optional[float]], Optional[float]] = None,
         rounding: Union[Sequence[Optional[str]], Optional[str]] = None,
         allow_missing_keys: bool = False,
+        **kwargs,
     ) -> None:
         """
         Args:
@@ -162,6 +167,9 @@ def __init__(
                 available options: ["torchrounding"]. it also can be a sequence of str or None,
                 each element corresponds to a key in ``keys``.
             allow_missing_keys: don't raise exception if key is missing.
+            kwargs: additional parameters to ``AsDiscrete``.
+                ``dim``, ``keepdim``, ``dtype`` are supported, unrecognized parameters will be ignored.
+                These default to ``0``, ``True``, ``torch.float`` respectively.
 
         """
         super().__init__(keys, allow_missing_keys)
@@ -180,6 +188,7 @@ def __init__(
 
         self.rounding = ensure_tuple_rep(rounding, len(self.keys))
         self.converter = AsDiscrete()
+        self.converter.kwargs = kwargs
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)


Test patch:
diff --git a/tests/test_activations.py b/tests/test_activations.py
--- a/tests/test_activations.py
+++ b/tests/test_activations.py
@@ -38,6 +38,15 @@
         ]
     )
 
+    TEST_CASES.append(
+        [
+            {"sigmoid": False, "softmax": True, "other": None, "unused": True, "dim": 1},
+            p([[[0.0, 1.0]], [[2.0, 3.0]]]),
+            p([[[1.0, 1.0]], [[1.0, 1.0]]]),
+            (2, 1, 2),
+        ]
+    )
+
     TEST_CASES.append(
         [
             {"sigmoid": False, "softmax": False, "other": torch.tanh},
diff --git a/tests/test_activationsd.py b/tests/test_activationsd.py
--- a/tests/test_activationsd.py
+++ b/tests/test_activationsd.py
@@ -21,7 +21,7 @@
 for p in TEST_NDARRAYS:
     TEST_CASES.append(
         [
-            {"keys": ["pred", "label"], "sigmoid": False, "softmax": [True, False], "other": None},
+            {"keys": ["pred", "label"], "sigmoid": False, "softmax": [True, False], "other": None, "dim": 0},
             {"pred": p([[[0.0, 1.0]], [[2.0, 3.0]]]), "label": p([[[0.0, 1.0]], [[2.0, 3.0]]])},
             {"pred": p([[[0.1192, 0.1192]], [[0.8808, 0.8808]]]), "label": p([[[0.0, 1.0]], [[2.0, 3.0]]])},
             (2, 1, 2),
diff --git a/tests/test_as_discrete.py b/tests/test_as_discrete.py
--- a/tests/test_as_discrete.py
+++ b/tests/test_as_discrete.py
@@ -29,7 +29,7 @@
 
     TEST_CASES.append(
         [
-            {"argmax": True, "to_onehot": 2, "threshold": 0.5},
+            {"argmax": True, "to_onehot": 2, "threshold": 0.5, "dim": 0},
             p([[[0.0, 1.0]], [[2.0, 3.0]]]),
             p([[[0.0, 0.0]], [[1.0, 1.0]]]),
             (2, 1, 2),
@@ -69,6 +69,11 @@ def test_value_shape(self, input_param, img, out, expected_shape):
         assert_allclose(result, out, rtol=1e-3, type_test="tensor")
         self.assertTupleEqual(result.shape, expected_shape)
 
+    def test_additional(self):
+        for p in TEST_NDARRAYS:
+            out = AsDiscrete(argmax=True, dim=1, keepdim=False)(p([[[0.0, 1.0]], [[2.0, 3.0]]]))
+            assert_allclose(out, p([[0.0, 0.0], [0.0, 0.0]]), type_test=False)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_as_discreted.py b/tests/test_as_discreted.py
--- a/tests/test_as_discreted.py
+++ b/tests/test_as_discreted.py
@@ -38,7 +38,7 @@
 
     TEST_CASES.append(
         [
-            {"keys": ["pred"], "argmax": True, "to_onehot": 2, "threshold": 0.5},
+            {"keys": ["pred"], "argmax": True, "to_onehot": 2, "threshold": 0.5, "dim": 0, "keepdim": True},
             {"pred": p([[[0.0, 1.0]], [[2.0, 3.0]]])},
             {"pred": p([[[0.0, 0.0]], [[1.0, 1.0]]])},
             (2, 1, 2),