Add EMA class for exponential moving average weights management

model.py

@@ -131,4 +131,38 @@ class Unet(torch.nn.Module):
         x7 = torch.cat([x2, x7], dim=1)
         x7 = self.resblock4(x7, t_emb)
 
-        return self.outc(x7)
+        return self.outc(x7)
+
+#EMA weights 
+class EMA():
+    def __init__(self, model, decay=0.9999):
+        self.model = model
+        self.decay = decay 
+        self.shadow = {}
+        self.backup = {}
+
+    def register(self):
+        for name , param in self.model.named_parameters():
+            if param.requires_grad:
+                self.shadow[name] = param.data.clone()
+
+    def upadate(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_gard:
+                assert name in self.shadow
+                new_avg = self.decay*self.shadow[name] + ((1. - self.decay)*param)
+                self.shadow[name] = new_avg.clone()
+
+    def apply_shadow(self):
+        for name , param in self.model.named_parameters():
+            if param.requires_grad:
+                assert name in self.shadow
+                self.backup = param.data.clone()
+                param.data.copy_(self.shadow[name])
+
+    def restore(self):
+        for name, param in self.model.named_parameters():
+            if param.require_grad:
+                assert name in self.backup 
+                param.data.copy_(self.backup[name])
+                self.backup = {}