mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix loss accumulator when buffers are not filled

Browse Source 
They were reporting incorrect losses.
This commit is contained in:
James Betker 2020-10-09 19:47:59 -06:00
parent fe50d6f9d0
commit 61e5047c60
1 changed files with 9 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
codes/utils/loss_accumulator.py
View File

@ -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