Enable compatibility with torchmetrics >= 0.6.0 (#27)

README.md

@@ -86,18 +86,18 @@ class. It's
 [**init**](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.trainer.trainer.html#pytorch_lightning.trainer.trainer.Trainer.__init__)
 method provides various configuration options.
 
-If you want to run K-Means with a GPU, you can pass the option `gpus=1` to the estimator's
-initializer:
+If you want to run K-Means with a GPU, you can pass the options `accelerator='gpu'` and `devices=1`
+to the estimator's initializer:
 
 ```python
-estimator = KMeans(3, trainer_params=dict(gpus=1))
+estimator = KMeans(3, trainer_params=dict(accelerator='gpu', devices=1))
 ```
 
 Similarly, if you want to train on 4 nodes simultaneously where each node has one GPU available,
 you can specify this as follows:
 
 ```python
-estimator = KMeans(3, trainer_params=dict(num_nodes=4, gpus=1))
+estimator = KMeans(3, trainer_params=dict(num_nodes=4, accelerator='gpu', devices=1))
 ```
 
 In fact, **you do not need to change anything else in your code**.

docs/index.rst

@@ -78,19 +78,19 @@ For GPU- and multi-node training, PyCave leverages PyTorch Lightning. The hardwa
 runs on is determined by the :class:`pytorch_lightning.trainer.Trainer` class. It's
 :meth:`~pytorch_lightning.trainer.Trainer.__init__` method provides various configuration options.
 
-If you want to run K-Means with a GPU, you can pass the option ``gpus=1`` to the estimator's
+If you want to run K-Means with a GPU, you can pass the option ``accelerator='gpu'`` and ``devices=1`` to the estimator's
 initializer:
 
 .. code-block:: python
 
-    estimator = KMeans(3, trainer_params=dict(gpus=1))
+    estimator = KMeans(3, trainer_params=dict(accelerator='gpu', devices=1))
 
 Similarly, if you want to train on 4 nodes simultaneously where each node has one GPU available,
 you can specify this as follows:
 
 .. code-block:: python
 
-    estimator = KMeans(3, trainer_params=dict(num_nodes=4, gpus=1))
+    estimator = KMeans(3, trainer_params=dict(num_nodes=4, accelerator='gpu', 1))
 
 In fact, **you do not need to change anything else in your code**.
 

pycave/bayes/gmm/lightning_module.py

@@ -2,7 +2,7 @@
 import pytorch_lightning as pl
 import torch
 from pytorch_lightning.callbacks import EarlyStopping
-from torchmetrics import AverageMeter
+from torchmetrics import MeanMetric
 from pycave.bayes.core import cholesky_precision
 from pycave.utils import NonparametricLightningModule
 from .metrics import CovarianceAggregator, MeanAggregator, PriorAggregator
@@ -65,7 +65,7 @@ def __init__(
         )
 
         # Initialize metrics
-        self.metric_nll = AverageMeter(dist_sync_fn=self.all_gather)
+        self.metric_nll = MeanMetric(dist_sync_fn=self.all_gather)
 
     def configure_callbacks(self) -> list[pl.Callback]:
         if self.convergence_tolerance == 0:

pycave/bayes/gmm/metrics.py

@@ -9,6 +9,8 @@ class PriorAggregator(Metric):
     The prior aggregator aggregates component probabilities over batches and process.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_components: int,
@@ -33,6 +35,8 @@ class MeanAggregator(Metric):
     The mean aggregator aggregates component means over batches and processes.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_components: int,
@@ -63,6 +67,8 @@ class CovarianceAggregator(Metric):
     The covariance aggregator aggregates component covariances over batches and processes.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_components: int,

pycave/bayes/markov_chain/lightning_module.py

@@ -1,6 +1,6 @@
 import torch
 from torch.nn.utils.rnn import PackedSequence
-from torchmetrics import AverageMeter
+from torchmetrics import MeanMetric
 from pycave.bayes.markov_chain.metrics import StateCountAggregator
 from pycave.utils import NonparametricLightningModule
 from .model import MarkovChainModel
@@ -27,7 +27,7 @@ def __init__(self, model: MarkovChainModel, symmetric: bool = False):
             symmetric=self.symmetric,
             dist_sync_fn=self.all_gather,
         )
-        self.metric_nll = AverageMeter(dist_sync_fn=self.all_gather)
+        self.metric_nll = MeanMetric(dist_sync_fn=self.all_gather)
 
     def on_train_epoch_start(self) -> None:
         self.aggregator.reset()

pycave/bayes/markov_chain/metrics.py

@@ -9,6 +9,8 @@ class StateCountAggregator(Metric):
     The state count aggregator aggregates initial states and transitions between states.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_states: int,

pycave/clustering/kmeans/lightning_module.py

@@ -4,7 +4,7 @@
 import pytorch_lightning as pl
 import torch
 from pytorch_lightning.callbacks import EarlyStopping
-from torchmetrics import AverageMeter
+from torchmetrics import MeanMetric
 from pycave.utils import NonparametricLightningModule
 from .metrics import (
     BatchAverager,
@@ -51,7 +51,7 @@ def __init__(
         )
 
         # Initialize metrics
-        self.metric_inertia = AverageMeter()
+        self.metric_inertia = MeanMetric()
 
     def configure_callbacks(self) -> List[pl.Callback]:
         if self.convergence_tolerance == 0:
@@ -239,8 +239,8 @@ def nonparametric_training_step(self, batch: torch.Tensor, batch_idx: int) -> No
 
     def nonparametric_training_epoch_end(self) -> None:
         if self.current_epoch == 0:
-            choice = self.uniform_sampler.compute()[0]
-            self.model.centroids[0].copy_(choice)
+            choice = self.uniform_sampler.compute()
+            self.model.centroids[0].copy_(choice[0] if choice.dim() > 0 else choice)
         elif self._is_current_epoch_sampling:
             candidates = self.distance_sampler.compute()
             self.centroid_candidates.copy_(candidates)

pycave/clustering/kmeans/metrics.py

@@ -9,6 +9,8 @@ class CentroidAggregator(Metric):
     The centroid aggregator aggregates kmeans centroids over batches and processes.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_clusters: int,
@@ -49,6 +51,8 @@ class UniformSampler(Metric):
     they were already sampled from).
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_choices: int,
@@ -109,6 +113,8 @@ class DistanceSampler(Metric):
     duplicates.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_choices: int,
@@ -169,6 +175,8 @@ class BatchSummer(Metric):
     Sums the values for a batch of items independently.
     """
 
+    full_state_update = True
+
     def __init__(self, num_values: int, *, dist_sync_fn: Optional[Callable[[Any], Any]] = None):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
@@ -187,6 +195,8 @@ class BatchAverager(Metric):
     Averages the values for a batch of items independently.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_values: int,

pyproject.toml

@@ -18,7 +18,7 @@ numpy = "^1.20.3"
 python = ">=3.8,<3.11"
 pytorch-lightning = "^1.6.0"
 torch = "^1.8.0"
-torchmetrics = "^0.5.1,<0.6.0"
+torchmetrics = "^0.6.0"
 
 [tool.poetry.group.pre-commit.dependencies]
 black = "^22.12.0"

tests/bayes/gmm/benchmark_gmm_estimator.py

@@ -128,6 +128,6 @@ def test_pycave_gpu(
         convergence_tolerance=0,
         covariance_regularization=1e-3,
         batch_size=batch_size,
-        trainer_params=dict(max_epochs=100, gpus=1),
+        trainer_params=dict(max_epochs=100, accelerator="gpu", devices=1),
     )
     benchmark(estimator.fit, data)

tests/clustering/kmeans/benchmark_kmeans_estimator.py

@@ -120,6 +120,6 @@ def test_pycave_gpu(
         init_strategy=init_strategy,
         batch_size=batch_size,
         convergence_tolerance=0,
-        trainer_params=dict(gpus=1, max_epochs=100),
+        trainer_params=dict(max_epochs=100, accelerator="gpu", devices=1),
     )
     benchmark(estimator.fit, data)