28 lines
578 B
Python
28 lines
578 B
Python
import os
|
|
import sys
|
|
import numpy as np
|
|
|
|
|
|
|
|
class Accumulator():
|
|
def __init__(self):
|
|
self.dict = dict()
|
|
|
|
def update(self, dict):
|
|
for key in dict.keys():
|
|
if key in self.dict.keys():
|
|
self.dict[key].append(dict[key].detach().cpu().numpy())
|
|
else:
|
|
self.dict[key] = [dict[key].detach().cpu().numpy()]
|
|
|
|
def get(self):
|
|
result = dict()
|
|
for key in self.dict.keys():
|
|
result[key] = np.mean(self.dict[key])
|
|
return result
|
|
|
|
def reset(self):
|
|
self.dict = dict()
|
|
|
|
|