From 61e5047c600fce8f1e127a89f7fd2f08080d9fd0 Mon Sep 17 00:00:00 2001 From: James Betker Date: Fri, 9 Oct 2020 19:47:59 -0600 Subject: [PATCH] Fix loss accumulator when buffers are not filled They were reporting incorrect losses. --- codes/utils/loss_accumulator.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/codes/utils/loss_accumulator.py b/codes/utils/loss_accumulator.py index a8762360..8965cf3d 100644 --- a/codes/utils/loss_accumulator.py +++ b/codes/utils/loss_accumulator.py @@ -8,17 +8,22 @@ class LossAccumulator: def add_loss(self, name, tensor): if name not in self.buffers.keys(): - self.buffers[name] = (0, torch.zeros(self.buffer_sz)) - i, buf = self.buffers[name] + self.buffers[name] = (0, torch.zeros(self.buffer_sz), False) + i, buf, filled = self.buffers[name] # Can take tensors or just plain python numbers. if isinstance(tensor, torch.Tensor): buf[i] = tensor.detach().cpu() else: buf[i] = tensor - self.buffers[name] = ((i+1) % self.buffer_sz, buf) + filled = i+1 >= self.buffer_sz or filled + self.buffers[name] = ((i+1) % self.buffer_sz, buf, filled) def as_dict(self): result = {} for k, v in self.buffers.items(): - result["loss_" + k] = torch.mean(v[1]) + i, buf, filled = v + if filled: + result["loss_" + k] = torch.mean(buf) + else: + result["loss_" + k] = torch.mean(buf[:i]) return result \ No newline at end of file