GROUPWORK_Fix SGD

p2/SGD.py

@@ -9,24 +9,18 @@ class SGD_Optim():
         self.velocity = []
 
     def step(self): 
-        parameters = self.model.param()
-        # print(parameters)
-        for p in parameters: # p = [(w, wgrad),(b, bgrad)] or [(None,None)]
+        modules = self.model.mods
+        
+        for module in modules: 
+            p = module.param()
             if (p[0][0] is not None):
-                (w, w_grad), (b, b_grad) = p
-                if not self.velocity:
-                    self.velocity.append((self.lr/(1 - self.gamma))*w_grad)
-                    self.velocity.append((self.lr/(1 - self.gamma))*b_grad)
-                v1 = self.gamma*self.velocity[0] + (self.lr/self.batch_size)*w_grad
-                v2 = self.gamma*self.velocity[1] + (self.lr/self.batch_size)*b_grad
-                self.velocity[0] = v1
-                self.velocity[1] = v2          
-                w -= v1
-                b -= v2
-                    
+                if module.velocity_w is None:
+                    module.velocity_w = ((self.lr/(1 - self.gamma))*module.grad_w)
+                if module.velocity_b is None:
+                    module.velocity_b =  ((self.lr/(1 - self.gamma))*module.grad_b)
+                module.velocity_w = self.gamma * module.velocity_w + (self.lr/self.batch_size)*module.grad_w
+                module.velocity_b = self.gamma * module.velocity_b + (self.lr/self.batch_size)*module.grad_b
+                module.update_params(opt=True, w=(module.w - module.velocity_w), b=(module.b - module.velocity_b))
+
     def zero_grad(self):
-        parameters = self.model.param()
-        for p in parameters: # p = [(w, wgrad),(b, bgrad)] or [(None,None)]
-            if p[0][1] is not None: # p[0] = (w, wgrad)
-                p[0][1].zero_() # p[0][1] = wgrad
-                p[1][1].zero_() # p[1][1] = bgrad
\ No newline at end of file
+        self.model.zero_grad()