diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py index 45c806c..1184a52 100644 --- a/src/cdtools/reconstructors/base.py +++ b/src/cdtools/reconstructors/base.py @@ -164,7 +164,8 @@ class Reconstructor: for inputs, patterns in self.data_loader: if hasattr(self.model, 'loss_normalizer') and \ self.model.loss_normalizer is not None: - self.model.loss_normalizer.accumulate(patterns) + self.model.loss_normalizer.accumulate( + patterns, mask=self.model.mask) N += 1 def closure(): diff --git a/src/cdtools/tools/losses/losses.py b/src/cdtools/tools/losses/losses.py index 5b304fc..bb65858 100644 --- a/src/cdtools/tools/losses/losses.py +++ b/src/cdtools/tools/losses/losses.py @@ -347,13 +347,13 @@ class SimplePoissonNLLNormalizer(object): """ if not self.first_pass_complete: + self.offset += poisson_nll(patterns, patterns, mask=mask) if mask is None: self.sum_nonzero += t.sum(patterns >= 1) - self.offset += poisson_nll(patterns, patterns) else: masked_pats = patterns.masked_select(mask) self.sum_nonzero += t.sum(masked_pats >= 1) - self.offset += poisson_nll(masked_pats, masked_pats) + def normalize_loss(self, loss): @@ -373,7 +373,7 @@ class SimplePoissonNLLNormalizer(object): if not self.first_pass_complete: self.normalization = 0.5 * self.sum_nonzero self.first_pass_complete = True - + return (loss - self.offset) / self.normalization