mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-22 13:30:11 +08:00
William Falcon
31b71483c4
* remove .item which causes sync issues * fixed gradient acc sched * fixed gradient acc sched
40 lines
1000 B
Python
40 lines
1000 B
Python
import torch
|
|
|
|
|
|
class TensorRunningMean(object):
|
|
"""
|
|
Tracks a running mean without graph references.
|
|
Round robbin for the mean
|
|
"""
|
|
def __init__(self, window_length):
|
|
self.window_length = window_length
|
|
self.reset()
|
|
self.last_idx = 0
|
|
|
|
def reset(self):
|
|
self.memory = torch.Tensor(self.window_length)
|
|
self.current_idx = 0
|
|
|
|
def last(self):
|
|
return self.memory[self.last_idx]
|
|
|
|
def append(self, x):
|
|
# map proper type for memory if they don't match
|
|
if self.memory.type() != x.type():
|
|
self.memory.type_as(x)
|
|
|
|
# store without grads
|
|
with torch.no_grad():
|
|
self.memory[self.current_idx] = x
|
|
self.last_idx = self.current_idx
|
|
|
|
# increase index
|
|
self.current_idx += 1
|
|
|
|
# reset index when hit limit of tensor
|
|
if self.current_idx >= self.window_length:
|
|
self.current_idx = 0
|
|
|
|
def mean(self):
|
|
return self.memory.mean()
|