started condition trainings

custom_utils/custom_utils.py

@@ -402,4 +402,26 @@ def group_inputs_for_levels(batch, config):
             uv_orig = batch['uv_posed' + res_modifier]
         levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
     conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
-    return levels_input, conds
+    return levels_input, conds
+
+
+
+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])
+            else:
+                self.dict[key] = [dict[key]]
+
+    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()