From d4edc92f8ee1f4042cd4e457b82541bc7866850b Mon Sep 17 00:00:00 2001
From: gzhangruipeng <gzhangruipeng@gmail.com>
Date: Sat, 14 Sep 2024 19:23:00 +0800
Subject: [PATCH] Update the model.py for the msmlp example

---
 examples/cnn_ms/msmlp/model.py | 32 +++++++++++++++++++++++++++-----
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/examples/cnn_ms/msmlp/model.py b/examples/cnn_ms/msmlp/model.py
index f6e2a2e72..49567160c 100644
--- a/examples/cnn_ms/msmlp/model.py
+++ b/examples/cnn_ms/msmlp/model.py
@@ -32,6 +32,7 @@
 
 singa_dtype = {"float16": tensor.float16, "float32": tensor.float32}
 
+
 #### self-defined loss begin
 
 ### from autograd.py
@@ -62,11 +63,13 @@ def backward(self, dy=1.0):
         dx *= dy
         return dx
 
+
 def se_loss(x):
     # assert x.shape == t.shape, "input and target shape different: %s, %s" % (
     #     x.shape, t.shape)
     return SumError()(x)[0]
 
+
 ### from layer.py
 class SumErrorLayer(Layer):
     """
@@ -79,6 +82,7 @@ def __init__(self):
     def forward(self, x):
         return se_loss(x)
 
+
 #### self-defined loss end
 
 class MSMLP(model.Model):
@@ -92,7 +96,6 @@ def __init__(self, data_size=10, perceptron_size=100, num_classes=10):
         self.linear1 = layer.Linear(perceptron_size)
         self.linear2 = layer.Linear(num_classes)
         self.softmax_cross_entropy = layer.SoftMaxCrossEntropy()
-
         self.sum_error = SumErrorLayer()
 
     def forward(self, inputs):
@@ -101,12 +104,24 @@ def forward(self, inputs):
         y = self.linear2(y)
         return y
 
-    def train_one_batch(self, x, y, synflow_flag, dist_option, spars):
+    def train_one_batch(self, x, y, dist_option, spars, synflow_flag):
+        # print ("in train_one_batch")
         out = self.forward(x)
-        loss = self.softmax_cross_entropy(out, y)
+        # print ("train_one_batch x.data: \n", x.data)
+        # print ("train_one_batch y.data: \n", y.data)
+        # print ("train_one_batch out.data: \n", out.data)
+        if synflow_flag:
+            # print ("sum_error")
+            loss = self.sum_error(out)
+        else:  # normal training
+            # print ("softmax_cross_entropy")
+            loss = self.softmax_cross_entropy(out, y)
+        # print ("train_one_batch loss.data: \n", loss.data)
 
         if dist_option == 'plain':
+            # print ("before pn_p_g_list = self.optimizer(loss)")
             pn_p_g_list = self.optimizer(loss)
+            # print ("after pn_p_g_list = self.optimizer(loss)")
         elif dist_option == 'half':
             self.optimizer.backward_and_update_half(loss)
         elif dist_option == 'partialUpdate':
@@ -119,7 +134,13 @@ def train_one_batch(self, x, y, synflow_flag, dist_option, spars):
             self.optimizer.backward_and_sparse_update(loss,
                                                       topK=False,
                                                       spars=spars)
+        # print ("len(pn_p_g_list): \n", len(pn_p_g_list))
+        # print ("len(pn_p_g_list[0]): \n", len(pn_p_g_list[0]))
+        # print ("pn_p_g_list[0][0]: \n", pn_p_g_list[0][0])
+        # print ("pn_p_g_list[0][1].data: \n", pn_p_g_list[0][1].data)
+        # print ("pn_p_g_list[0][2].data: \n", pn_p_g_list[0][2].data)
         return pn_p_g_list, out, loss
+        # return pn_p_g_list[0], pn_p_g_list[1], pn_p_g_list[2], out, loss
 
     def set_optimizer(self, optimizer):
         self.optimizer = optimizer
@@ -127,9 +148,10 @@ def set_optimizer(self, optimizer):
 
 def create_model(pretrained=False, **kwargs):
     """Constructs a CNN model.
+
     Args:
         pretrained (bool): If True, returns a pre-trained model.
-    
+
     Returns:
         The created CNN model.
     """
@@ -196,4 +218,4 @@ def create_model(pretrained=False, **kwargs):
         out, loss = model(tx, ty, 'fp32', spars=None)
 
         if i % 100 == 0:
-            print("training loss = ", tensor.to_numpy(loss)[0])
+            print("training loss = ", tensor.to_numpy(loss)[0])
\ No newline at end of file