Refactor/move args to kwargs

CHANGELOG.md

@@ -23,12 +23,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `ClasswiseWrapper` for better logging of classification metrics with multiple output values ([#832](https://github.com/PyTorchLightning/metrics/pull/832))
 
 
+- Added `**kwargs` argument for passing additional arguments to base class ([#833](https://github.com/PyTorchLightning/metrics/pull/833))
+
+
 ### Changed
 
 
 ### Deprecated
 
-- Deprecated method `compute_on_step` ([#792](https://github.com/PyTorchLightning/metrics/pull/792))
+- Deprecated argument `compute_on_step` ([#792](https://github.com/PyTorchLightning/metrics/pull/792))
+
+
+- Deprecated passing in `dist_sync_on_step`, `process_group`, `dist_sync_fn` direct argument ([#833](https://github.com/PyTorchLightning/metrics/pull/833))
 
 
 ### Removed

docs/source/pages/overview.rst

@@ -348,3 +348,24 @@ In practise this means that:
     val = metric.compute() # this value cannot be back-propagated
 
 A functional metric is differentiable if its corresponding modular metric is differentiable.
+
+.. _Metric kwargs:
+
+*****************************
+Advanced distributed settings
+*****************************
+
+If you are running in a distributed environment, ``TorchMetrics`` will automatically take care of the distributed
+synchronization for you. However, the following three keyword arguments can be given to any metric class for
+further control over the distributed aggregation:
+
+- ``dist_sync_on_step``: This argument is ``bool`` that indicates if the metric should syncronize between
+  different devices every time ``forward`` is called. Setting this to ``True`` is in general not recommended
+  as syncronization is an expensive operation to do after each batch.
+
+- ``process_group``: By default we syncronize across the *world* i.e. all proceses being computed on. You
+  can provide an ``torch._C._distributed_c10d.ProcessGroup`` in this argument to specify exactly what
+  devices should be syncronized over.
+
+- ``dist_sync_fn``: By default we use :func:`torch.distributed.all_gather` to perform the synchronization between
+  devices. Provide another callable function for this argument to perform custom distributed synchronization.

tests/test_deprecated.py

@@ -0,0 +1,10 @@
+import pytest
+
+from torchmetrics import Accuracy
+
+
+def test_compute_on_step():
+    with pytest.warns(
+        DeprecationWarning, match="Argument `compute_on_step` is deprecated in v0.8 and will be removed in v0.9"
+    ):
+        Accuracy(compute_on_step=False)  # any metric will raise the warning

torchmetrics/audio/pesq.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Optional
+from typing import Any, Dict, Optional
 
 from torch import Tensor, tensor
 
@@ -49,15 +49,8 @@ class PerceptualEvaluationSpeechQuality(Metric):
             .. deprecated:: v0.8
                 Argument has no use anymore and will be removed v0.9.
 
-        dist_sync_on_step:
-            Synchronize metric state across processes at each ``forward()``
-            before returning the value at the step
-        process_group:
-            Specify the process group on which synchronization is called.
-            default: ``None`` (which selects the entire world)
-        dist_sync_fn:
-            Callback that performs the allgather operation on the metric state. When ``None``, DDP
-            will be used to perform the allgather
+        kwargs:
+            Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
     Raises:
         ModuleNotFoundError:
@@ -94,16 +87,9 @@ def __init__(
         fs: int,
         mode: str,
         compute_on_step: Optional[bool] = None,
-        dist_sync_on_step: bool = False,
-        process_group: Optional[Any] = None,
-        dist_sync_fn: Optional[Callable[[Tensor], Tensor]] = None,
+        **kwargs: Dict[str, Any],
     ) -> None:
-        super().__init__(
-            compute_on_step=compute_on_step,
-            dist_sync_on_step=dist_sync_on_step,
-            process_group=process_group,
-            dist_sync_fn=dist_sync_fn,
-        )
+        super().__init__(compute_on_step=compute_on_step, **kwargs)
         if not _PESQ_AVAILABLE:
             raise ModuleNotFoundError(
                 "PerceptualEvaluationSpeechQuality metric requires that `pesq` is installed."

torchmetrics/audio/pit.py

@@ -43,16 +43,9 @@ class PermutationInvariantTraining(Metric):
             .. deprecated:: v0.8
                 Argument has no use anymore and will be removed v0.9.
 
-        dist_sync_on_step:
-            Synchronize metric state across processes at each ``forward()``
-            before returning the value at the step.
-        process_group:
-            Specify the process group on which synchronization is called.
-        dist_sync_fn:
-            Callback that performs the allgather operation on the metric state. When `None`, DDP
-            will be used to perform the allgather.
         kwargs:
-            additional args for metric_func
+            Additional keyword arguments for either the `metric_func` or distributed communication,
+            see :ref:`Metric kwargs` for more info.
 
     Returns:
         average PermutationInvariantTraining metric
@@ -83,17 +76,9 @@ def __init__(
         metric_func: Callable,
         eval_func: str = "max",
         compute_on_step: Optional[bool] = None,
-        dist_sync_on_step: bool = False,
-        process_group: Optional[Any] = None,
-        dist_sync_fn: Optional[Callable[[Tensor], Tensor]] = None,
         **kwargs: Dict[str, Any],
     ) -> None:
-        super().__init__(
-            compute_on_step=compute_on_step,
-            dist_sync_on_step=dist_sync_on_step,
-            process_group=process_group,
-            dist_sync_fn=dist_sync_fn,
-        )
+        super().__init__(compute_on_step=compute_on_step, **kwargs)
         self.metric_func = metric_func
         self.eval_func = eval_func
         self.kwargs = kwargs

torchmetrics/audio/sdr.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Dict, Optional
 
 from torch import Tensor, tensor
 
@@ -53,14 +53,8 @@ class SignalDistortionRatio(Metric):
             .. deprecated:: v0.8
                 Argument has no use anymore and will be removed v0.9.
 
-        dist_sync_on_step:
-            Synchronize metric state across processes at each ``forward()``
-            before returning the value at the step.
-        process_group:
-            Specify the process group on which synchronization is called.
-        dist_sync_fn:
-            Callback that performs the allgather operation on the metric state. When `None`, DDP
-            will be used to perform the allgather.
+        kwargs:
+            Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
     Raises:
         ModuleNotFoundError:
@@ -115,21 +109,14 @@ def __init__(
         zero_mean: bool = False,
         load_diag: Optional[float] = None,
         compute_on_step: Optional[bool] = None,
-        dist_sync_on_step: bool = False,
-        process_group: Optional[Any] = None,
-        dist_sync_fn: Optional[Callable[[Tensor], Tensor]] = None,
+        **kwargs: Dict[str, Any],
     ) -> None:
         if not _FAST_BSS_EVAL_AVAILABLE:
             raise ModuleNotFoundError(
                 "SDR metric requires that `fast-bss-eval` is installed."
                 " Either install as `pip install torchmetrics[audio]` or `pip install fast-bss-eval`."
             )
-        super().__init__(
-            compute_on_step=compute_on_step,
-            dist_sync_on_step=dist_sync_on_step,
-            process_group=process_group,
-            dist_sync_fn=dist_sync_fn,
-        )
+        super().__init__(compute_on_step=compute_on_step, **kwargs)
 
         self.use_cg_iter = use_cg_iter
         self.filter_length = filter_length

torchmetrics/audio/snr.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Dict, Optional
 
 from torch import Tensor, tensor
 
@@ -43,14 +43,8 @@ class SignalNoiseRatio(Metric):
             .. deprecated:: v0.8
                 Argument has no use anymore and will be removed v0.9.
 
-        dist_sync_on_step:
-            Synchronize metric state across processes at each ``forward()``
-            before returning the value at the step.
-        process_group:
-            Specify the process group on which synchronization is called.
-        dist_sync_fn:
-            Callback that performs the allgather operation on the metric state. When `None`, DDP
-            will be used to perform the allgather.
+        kwargs:
+            Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
     Raises:
         TypeError:
@@ -82,16 +76,9 @@ def __init__(
         self,
         zero_mean: bool = False,
         compute_on_step: Optional[bool] = None,
-        dist_sync_on_step: bool = False,
-        process_group: Optional[Any] = None,
-        dist_sync_fn: Optional[Callable[[Tensor], Tensor]] = None,
+        **kwargs: Dict[str, Any],
     ) -> None:
-        super().__init__(
-            compute_on_step=compute_on_step,
-            dist_sync_on_step=dist_sync_on_step,
-            process_group=process_group,
-            dist_sync_fn=dist_sync_fn,
-        )
+        super().__init__(compute_on_step=compute_on_step, **kwargs)
         self.zero_mean = zero_mean
 
         self.add_state("sum_snr", default=tensor(0.0), dist_reduce_fx="sum")

torchmetrics/audio/stoi.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Optional
+from typing import Any, Dict, Optional
 
 from torch import Tensor, tensor
 
@@ -52,14 +52,8 @@ class ShortTimeObjectiveIntelligibility(Metric):
             .. deprecated:: v0.8
                 Argument has no use anymore and will be removed v0.9.
 
-        dist_sync_on_step:
-            Synchronize metric state across processes at each ``forward()``
-            before returning the value at the step.
-        process_group:
-            Specify the process group on which synchronization is called.
-        dist_sync_fn:
-            Callback that performs the allgather operation on the metric state. When `None`, DDP
-            will be used to perform the allgather.
+        kwargs:
+            Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
     Returns:
         average STOI value
@@ -101,16 +95,9 @@ def __init__(
         fs: int,
         extended: bool = False,
         compute_on_step: Optional[bool] = None,
-        dist_sync_on_step: bool = False,
-        process_group: Optional[Any] = None,
-        dist_sync_fn: Optional[Callable[[Tensor], Tensor]] = None,
+        **kwargs: Dict[str, Any],
     ) -> None:
-        super().__init__(
-            compute_on_step=compute_on_step,
-            dist_sync_on_step=dist_sync_on_step,
-            process_group=process_group,
-            dist_sync_fn=dist_sync_fn,
-        )
+        super().__init__(compute_on_step=compute_on_step, **kwargs)
         if not _PYSTOI_AVAILABLE:
             raise ModuleNotFoundError(
                 "STOI metric requires that `pystoi` is installed."

torchmetrics/classification/accuracy.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Optional
+from typing import Any, Dict, Optional
 
 from torch import Tensor, tensor
 
@@ -132,15 +132,8 @@ class Accuracy(StatScores):
             .. deprecated:: v0.8
                 Argument has no use anymore and will be removed v0.9.
 
-        dist_sync_on_step:
-            Synchronize metric state across processes at each ``forward()``
-            before returning the value at the step
-        process_group:
-            Specify the process group on which synchronization is called.
-            default: ``None`` (which selects the entire world)
-        dist_sync_fn:
-            Callback that performs the allgather operation on the metric state. When ``None``, DDP
-            will be used to perform the allgather
+        kwargs:
+            Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
     Raises:
         ValueError:
@@ -184,9 +177,7 @@ def __init__(
         multiclass: Optional[bool] = None,
         subset_accuracy: bool = False,
         compute_on_step: Optional[bool] = None,
-        dist_sync_on_step: bool = False,
-        process_group: Optional[Any] = None,
-        dist_sync_fn: Callable = None,
+        **kwargs: Dict[str, Any],
     ) -> None:
         allowed_average = ["micro", "macro", "weighted", "samples", "none", None]
         if average not in allowed_average:
@@ -201,9 +192,7 @@ def __init__(
             multiclass=multiclass,
             ignore_index=ignore_index,
             compute_on_step=compute_on_step,
-            dist_sync_on_step=dist_sync_on_step,
-            process_group=process_group,
-            dist_sync_fn=dist_sync_fn,
+            **kwargs,
         )
 
         if top_k is not None and (not isinstance(top_k, int) or top_k <= 0):

torchmetrics/classification/auc.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, List, Optional
+from typing import Any, Dict, List, Optional
 
 from torch import Tensor
 
@@ -38,14 +38,8 @@ class AUC(Metric):
             .. deprecated:: v0.8
                 Argument has no use anymore and will be removed v0.9.
 
-        dist_sync_on_step:
-            Synchronize metric state across processes at each ``forward()``
-            before returning the value at the step.
-        process_group:
-            Specify the process group on which synchronization is called.
-        dist_sync_fn:
-            Callback that performs the ``allgather`` operation on the metric state. When ``None``, DDP
-            will be used to perform the ``allgather``.
+        kwargs:
+            Additional keyword arguments, see :ref:`Metric kwargs` for more info.
     """
     is_differentiable = False
     x: List[Tensor]
@@ -55,16 +49,9 @@ def __init__(
         self,
         reorder: bool = False,
         compute_on_step: Optional[bool] = None,
-        dist_sync_on_step: bool = False,
-        process_group: Optional[Any] = None,
-        dist_sync_fn: Callable = None,
+        **kwargs: Dict[str, Any],
     ) -> None:
-        super().__init__(
-            compute_on_step=compute_on_step,
-            dist_sync_on_step=dist_sync_on_step,
-            process_group=process_group,
-            dist_sync_fn=dist_sync_fn,
-        )
+        super().__init__(compute_on_step=compute_on_step, **kwargs)
 
         self.reorder = reorder