Repository: intel/ipex-llm
Pull request: [Issue 621]Support set model in optimizer
## What changes were proposed in this pull request?
1. add set model in optimizer
2. currently optimizer takes a long time to create and initialize, to avoid confusion, explicitly tell user it's loading data

## How was this patch tested?
unit test and manual test

## Related links or issues (optional)
fixed https://github.com/intel-analytics/BigDL/issues/621



Problem statement:
refactor the optimizer initialization process in python api
1.  add *set_model()* to optimizer 
Currently the model could only be set to optimizer during creation. That means we can't reuse an optimizer for different models.  Creating optimizer may be time consuming sometimes. 
1. add *prepare_input()* to optimizer
Currently optimizer takes a long time to create and initialize.  The reason is loading input and caching takes time. We could warp this part of work into one function  and name it something like prepare_input/load_input, and allow user to train different models on the same set of input data. and it won't be confusing why initialization takes so long.

Hints from issue discussion:

Commit message: support set model in optimizer

Patch to review:
diff --git a/pyspark/dl/optim/optimizer.py b/pyspark/dl/optim/optimizer.py
index 4a769e5edb3..b282b8f77d8 100644
--- a/pyspark/dl/optim/optimizer.py
+++ b/pyspark/dl/optim/optimizer.py
@@ -188,6 +188,14 @@ def setvalidation(self, batch_size, val_rdd, trigger, val_method=["Top1Accuracy"
         callBigDlFunc(self.bigdl_type, "setValidation", self.value, batch_size,
                       trigger, val_rdd, val_method)
 
+    def setmodel(self, model):
+        """
+        Set model. 
+
+        :param model: new model
+        """
+        self.value.setModel(model.value)
+
     def setcheckpoint(self, checkpoint_trigger,
                       checkpoint_path, isOverWrite=True):
         """
diff --git a/pyspark/test/simple_integration_test.py b/pyspark/test/simple_integration_test.py
index 7a12262e4cb..aa2dda2e04b 100644
--- a/pyspark/test/simple_integration_test.py
+++ b/pyspark/test/simple_integration_test.py
@@ -55,6 +55,11 @@ def gen_rand_sample():
         trainingData = self.sc.parallelize(range(0, data_len)).map(
             lambda i: gen_rand_sample())
 
+        model_test = Sequential()
+        l1_test = Linear(3, 1, "Xavier").set_name("linear1_test")
+        self.assertEqual("linear1_test", l1_test.name())
+        model_test.add(l1_test)
+
         model = Sequential()
         l1 = Linear(FEATURES_DIM, 1, "Xavier").set_name("linear1")
         self.assertEqual("linear1", l1.name())
@@ -64,7 +69,7 @@ def gen_rand_sample():
                  "learningRateDecay": 0.0002,
                  "learingRateSchedule": Poly(0.5, int((data_len/batch_size)*epoch_num))}  # noqa
         optimizer = Optimizer(
-            model=model,
+            model=model_test,
             training_rdd=trainingData,
             criterion=MSECriterion(),
             optim_method="sgd",
@@ -77,6 +82,10 @@ def gen_rand_sample():
             trigger=EveryEpoch(),
             val_method=["Top1Accuracy"]
         )
+
+        optimizer.optimize()
+
+        optimizer.setmodel(model=model)
         tmp_dir = tempfile.mkdtemp()
         optimizer.setcheckpoint(SeveralIteration(1), tmp_dir)
         train_summary = TrainSummary(log_dir=tmp_dir,
@@ -88,12 +97,12 @@ def gen_rand_sample():
         optimizer.set_val_summary(val_summary)
 
         trained_model = optimizer.optimize()
-
         lr_result = train_summary.read_scalar("LearningRate")
         top1_result = val_summary.read_scalar("Top1Accuracy")
 
         # TODO: add result validation
         parameters = trained_model.parameters()
+
         self.assertIsNotNone(parameters["linear1"])
         print("parameters %s" % parameters["linear1"])
         predict_result = trained_model.predict(trainingData)
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/dataset/DataSet.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/dataset/DataSet.scala
index ae9889039ab..42b15fb3408 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/dataset/DataSet.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/dataset/DataSet.scala
@@ -172,7 +172,6 @@ trait DistributedDataSet[T] extends AbstractDataSet[T, RDD[T]] {
       preDataSet.originRDD().mapPartitions(_ => Iterator
         .single(broadcast.value.cloneTransformer())
       ).setName("Cached Transformer").persist()
-    cachedTransformer.count()
 
     new DistributedDataSet[C] {
       override def size(): Long = preDataSet.size()
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizer.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizer.scala
index 5539d94f491..bc0db735dbe 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizer.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizer.scala
@@ -677,6 +677,10 @@ class DistriOptimizer[T: ClassTag] (
 
   private var models: RDD[DistriOptimizer.Cache[T]] = null
 
+  import DistriOptimizer._
+  logger.info("Loading and caching data ...")
+  dataset.originRDD().cache().count()
+
   /**
    * Clean some internal states, so this or other optimizers can run optimize again
    *
@@ -702,10 +706,10 @@ class DistriOptimizer[T: ClassTag] (
     val coresPerNode = Engine.coreNumber()
 
     val partitionNum = dataset.originRDD().partitions.length
-    val size = model.getParameters()._1.nElement()
+    val size = _model.getParameters()._1.nElement()
     val parameters = AllReduceParameter.newParameter(partitionNum, size)
 
-    models = DistriOptimizer.initThreadModels(model, dataset, criterion, state,
+    models = DistriOptimizer.initThreadModels(_model, dataset, criterion, state,
       nodeNumber, coresPerNode, checkSingleton, parameters, validationMethods)
 
     if (checkpointPath.isDefined) {
@@ -766,7 +770,7 @@ class DistriOptimizer[T: ClassTag] (
               state = T.load(stateFile)
               DistriOptimizer.logger.info("Recover from last snapshot")
             } else {
-              newModel = model
+              newModel = _model
               DistriOptimizer.logger.info("Recover from origin model")
             }
             optimMethod.clearHistory(state)
@@ -782,12 +786,12 @@ class DistriOptimizer[T: ClassTag] (
 
     val trainedModel = DistriOptimizer.getModel(models, parameters)
 
-    nn.Utils.copyModule(trainedModel, model)
+    nn.Utils.copyModule(trainedModel, _model)
 
     // Reset some internal states, so this or other optimizers can run optimize again
     clearState()
 
-    model
+    _model
   }
 
   private def getLatestFile(path: String, fileName: String): String = {
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/Optimizer.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/Optimizer.scala
index 9e517bde379..1d60e6227c9 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/Optimizer.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/Optimizer.scala
@@ -44,6 +44,7 @@ abstract class Optimizer[T: ClassTag, D](
   protected val dataset: DataSet[D],
     protected val criterion: Criterion[T])(implicit ev : TensorNumeric[T])
 {
+  protected var _model: Module[T] = model
   protected var state: Table = T()
   protected var optimMethod: OptimMethod[T] = new SGD[T]()
   protected var endWhen: Trigger = Trigger.maxIteration(100)
@@ -172,6 +173,16 @@ abstract class Optimizer[T: ClassTag, D](
     this
   }
 
+  /**
+   * Set a model to the optimizer
+   *
+   * @param model new model
+   */
+  def setModel(model: Module[T]): this.type = {
+    this._model = model
+    this
+  }
+
   /**
    * Set a state(learning rate, epochs...) to the optimizer
    *
diff --git a/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/DistriOptimizerSpec.scala b/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/DistriOptimizerSpec.scala
index 9f4228c02c6..2a202225602 100644
--- a/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/DistriOptimizerSpec.scala
+++ b/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/DistriOptimizerSpec.scala
@@ -169,6 +169,25 @@ class DistriOptimizerSpec extends FlatSpec with Matchers with BeforeAndAfter {
     result2(Array(1)) should be(1.0 +- 5e-2)
   }
 
+  "Train with MSE and SGD" should "be trained with good result after reset model" in {
+    var mm = cre
+    val optimizer = new DistriOptimizer[Double](mm, dataSet, new MSECriterion[Double]())
+      .setState(T("learningRate" -> 20.0))
+      .setEndWhen(Trigger.maxEpoch(5))
+    optimizer.optimize()
+
+    mm = mse
+    mm.getParameters()._1.fill(0.125)
+    optimizer.setModel(mm)
+    val model = optimizer.optimize()
+
+    val result1 = model.forward(input1).asInstanceOf[Tensor[Double]]
+    result1(Array(1)) should be(0.0 +- 5e-2)
+
+    val result2 = model.forward(input2).asInstanceOf[Tensor[Double]]
+    result2(Array(1)) should be(1.0 +- 5e-2)
+  }
+
   it should "be same compare to ref optimizer" in {
     RandomGenerator.RNG.setSeed(10)
     val optimizer = new DistriOptimizer(


Code review comments:
diff_hunk: @@ -677,6 +677,12 @@ class DistriOptimizer[T: ClassTag] (
 
   private var models: RDD[DistriOptimizer.Cache[T]] = null
 
+  import DistriOptimizer._
+  if (dataset.originRDD != null) {
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizer.scala
text: @user1:
From the issue, wrap this code in a method?

diff_hunk: @@ -677,6 +677,12 @@ class DistriOptimizer[T: ClassTag] (
 
   private var models: RDD[DistriOptimizer.Cache[T]] = null
 
+  import DistriOptimizer._
+  if (dataset.originRDD != null) {
+    logger.info("Loading and caching data ...")
+    dataset.originRDD().cache().count()
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizer.scala
text: @user1:
The rdd has been already cached, why cache it again?

@author:
My origin thought is to avoid user created his own rdd and didn't cache it. Since the optimizer use dataset.originRDD to training, what about we cache dataset.originRDD here and remove cache operation in dataset.scala.

@user1:
I don't agree optimizer do anything related with manipulate data, which should be delegate to dataset. The name of dataset is CachedDataset, so it cache the data for performance. If we have a new dataset named NoCacheDataSet to handle extra large data, will you cache it again?

diff_hunk: @@ -188,6 +188,14 @@ def setvalidation(self, batch_size, val_rdd, trigger, val_method=["Top1Accuracy"
         callBigDlFunc(self.bigdl_type, "setValidation", self.value, batch_size,
                       trigger, val_rdd, val_method)
 
+    def setmodel(self, model):
path: pyspark/dl/optim/optimizer.py
text: @user1:
From the issue, I think you should add another prepare_input method. And in python, the naming style should be set_model instead of setmodel

@author:
I checked other methods in optimizer.py, looks like the naming style is Java styple(eg, setvalidation, setcheckpoint). I use the same style to make sure the consistent style.

@author:
I didn't add a separate method perpare_input as I think the concern here is when create optimizer, it silently load/caching data which makes user confuse why create optimizer takes so long. I think it should be fine if we explicit told user we are loading data. Besides I think prepare_input is used to tell user we are loading data. I am not very likely to add the method and ask user call it. But if most of you think it's necessary to add the method, I will do it.


@user1:
but you can see there are set_train_summary, set_val_summary in the same class. And most python classes use the under_score style. Why don't you follow the majority? Please use the underscore style, and change the other two:)

@user2-li looks like code styled in python is not consisted.

@user1:
Again, I think optimizer should not cache the data. It's not a job of the optimizer.

I'm agreed with your opinion. A prepare_input is misleading. I think the problem mentioned in the issue only a problem in some interactive programming(e.g. notebook), user may feel it takes long time when execute create optimizer. It's not a problem in a batch execution.

A prepare_input allow notebook user to seperate the data loading time and training time. If you don't have it, user may think the training take so long time, in which most is spent on loading data. I think maybe we provide this method in optimizer, even user didn't call that method, optimizer will do it at the start of training. And don't cache the data directly in optimizer. A dataset should also provide a method prepareData to hide the details.