Allow GMM evaluation to return NLL for individual datapoints

Author: borchero

pycave/bayes/_internal/utils.py

@@ -80,7 +80,7 @@ def log_responsibilities(log_probs, comp_priors, return_log_likelihood=False):
     log_resp = posterior - evidence
 
     if return_log_likelihood:
-        return log_resp, evidence.sum()
+        return log_resp, evidence
     return log_resp
 
 

pycave/bayes/gmm/engine.py

@@ -49,14 +49,22 @@ def after_epoch(self, _):
     def after_training(self):
         self.cache = None
 
-    def evaluate(self, data, metrics=None, **kwargs):
-        metrics = {'neg_log_likelihood': lambda x: x}
-        return super().evaluate(data, metrics=metrics, **kwargs)
+    def evaluate(self, data, reduction='mean', **kwargs):
+        def reduce(x):
+            if reduction == 'mean':
+                return x.mean()
+            if reduction == 'sum':
+                return x.sum()
+            return x
+
+        metrics = {'neg_log_likelihood': reduce}
+        result = super().evaluate(data, metrics=metrics, **kwargs)
+        return result['neg_log_likelihood']
 
     def train_batch(self, data, eps=0.01):
         # E-step: compute responsibilities
         responsibilities, nll = self.model(data)
-        nll_ = nll.item() / data.size(0)
+        nll_ = nll.mean().item()
 
         # M-step: maximize
         gaussian_max = self.model.gaussian.maximize(data, responsibilities, self.requires_batching)
@@ -78,10 +86,7 @@ def train_batch(self, data, eps=0.01):
         self.cache['eps'] = eps
 
     def eval_batch(self, data):
-        return {
-            'nll': self.model(data)[1].item(), # nll
-            'n': data.data.size(0)
-        }
+        return self.model(data)[1] # NLL for all data samples
 
     def predict_batch(self, data):
         # Get responsibilities and normalize them to get a distribution over components
@@ -90,9 +95,3 @@ def predict_batch(self, data):
     def collate_losses(self, _):
         nll = self.cache['neg_log_likelihood']
         return {'neg_log_likelihood': nll}
-
-    def collate_evals(self, evals):
-        # Only negative log-likelihood
-        nll_sum = sum([p['nll'] for p in evals])
-        n = sum([p['n'] for p in evals])
-        return torch.as_tensor(nll_sum / n)

pycave/bayes/gmm/model.py

@@ -38,6 +38,11 @@ class may be used to find clusters whenever you expect data to be generated from
         eps: float, default: 0.01
             The minimum per-datapoint difference in the negative log-likelihood to consider a
             model "better", thus indicating convergence.
+
+    `evaluate(...)`
+        reduction: str, default: 'mean'
+            The reduction performed for the negative log-likelihood as for common PyTorch metrics.
+            Must be one of ['mean', 'sum', 'none'].
     """
 
     __engine__ = GMMEngine
@@ -118,14 +123,14 @@ def forward(self, data):
         -------
         torch.Tensor [N, K]
             The responsibilities for each datapoint and component (number of components K).
-        torch.Tensor [1]
-            The negative log-likelihood of the data.
+        torch.Tensor [N]
+            The negative log-likelihood for all data samples.
         """
         probs = self.gaussian.evaluate(data, log=True)
         log_resp, log_likeli = log_responsibilities(
             probs, self.component_weights, return_log_likelihood=True
         )
-        return log_resp.exp(), -log_likeli
+        return log_resp.exp(), -log_likeli.squeeze(-1)
 
     def sample(self, n, return_components=False):
         """