Repository: intel/ipex-llm
Pull request: [Issue 864] save module weights, load (all / part) module weights
## What changes were proposed in this pull request?

load and save weights.

This is useful for finetuning, If you need to load weights into a different architecture (with some layers in common), for instance for fine-tuning or transfer-learning, you can load weights by layer name.

## How was this patch tested?

unit test

## Related links or issues (optional)
fixed https://github.com/intel-analytics/BigDL/issues/864



Problem statement:
Load part of parameters from saved bigdl model
Now BigDL can load serialized BigDL model (without the model definition).
But it does not support loading part of parameters to an existing defined model.

Use case: Finetune a model. 
Since many models are extension of some classic model, e.g. VGG16. As I want to finetune the whole model with pretrained base model, we need to only load parameters of pretrained base model.

Hints from issue discussion:
In keras, they can save_weights and load_weights

```
model.save_weights('my_model_weights.h5')
model.load_weights('my_model_weights.h5')
```
If you need to load weights into a different architecture (with some layers in common), for instance for fine-tuning or transfer-learning, you can load weights by layer name:

model.load_weights('my_model_weights.h5', by_name=True)

I think it is good to support functions like the above

Commit message: add support for loading from another model

Patch to review:
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala
index 836aef196d9..ba00d15bbd1 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala
@@ -431,6 +431,75 @@ abstract class AbstractModule[A <: Activity: ClassTag, B <: Activity: ClassTag,
     }
   }
 
+  /**
+   * save weights and bias to file
+   * @param path file to save
+   * @param overWrite whether to overwrite or not
+   */
+  def saveWeights(path: String, overWrite: Boolean): Unit = {
+    val parameterTable = getParametersTable()
+    val weightsBiasTable = T()
+    parameterTable.foreach {
+      case (name: String, params: Table) =>
+        val wb = T()
+        if (params.contains("weight")) {
+          wb("weight") = params("weight")
+        }
+        if (params.contains("bias")) {
+          wb("bias") = params("bias")
+        }
+        weightsBiasTable(name) = wb
+      case _ => throw new UnsupportedOperationException("invalid parameter table")
+    }
+    weightsBiasTable.save(path, overWrite)
+  }
+
+  /**
+   * load pretrained weights and bias to current module
+   * @param weightPath file to store weights and bias
+   * @param matchAll whether to match all layers' weights and bias,
+   *                 if not, only load existing pretrained weights and bias
+   * @return current module
+   */
+  def loadWeights(weightPath: String, matchAll: Boolean = true): this.type = {
+    val srcParameter = File.load[Table](weightPath)
+    val targetParameter = getParametersTable()
+    copyWeights(targetParameter, srcParameter)
+    this
+  }
+
+  /**
+   * copy weights from another model, mapping by layer name
+   * @param srcModel model to copy from
+   * @param matchAll whether to match all layers' weights and bias,
+   * @return current module
+   */
+  def loadModelWeights(srcModel: Module[Float], matchAll: Boolean = true): this.type = {
+    val srcParameters = srcModel.getParametersTable()
+    val parameterTable = getParametersTable()
+    copyWeights(parameterTable, srcParameters, matchAll)
+    this
+  }
+
+  private def copyWeights(target: Table, src: Table, matchAll: Boolean = true): Unit = {
+    target.foreach {
+      case (name: String, targetParams: Table) =>
+        if (src.contains(name)) {
+          val srcParams = src[Table](name)
+          if (srcParams.contains("weight")) {
+            val w = srcParams[Tensor[T]]("weight")
+            targetParams[Tensor[T]]("weight").resizeAs(w).copy(w)
+          }
+          if (srcParams.contains("bias")) {
+            val b = srcParams[Tensor[T]]("bias")
+            targetParams[Tensor[T]]("bias").resizeAs(b).copy(b)
+          }
+        } else {
+          if (matchAll) new Exception(s"module $name cannot find corresponding weight bias")
+        }
+    }
+  }
+
   /**
    * Some other modules point to current module
    * @param nodes upstream module nodes
diff --git a/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/AbstractModuleSpec.scala b/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/AbstractModuleSpec.scala
index ac0e751e662..bf330adcd37 100644
--- a/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/AbstractModuleSpec.scala
+++ b/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/AbstractModuleSpec.scala
@@ -15,6 +15,7 @@
  */
 package com.intel.analytics.bigdl.nn
 
+import com.intel.analytics.bigdl._
 import org.scalatest.{FlatSpec, Matchers}
 import com.intel.analytics.bigdl.numeric.NumericFloat
 
@@ -71,4 +72,137 @@ class AbstractModuleSpec extends FlatSpec with Matchers {
       s("module").get
     }
   }
+
+  "weights save and load" should "work properly" in {
+    val tmpFile = java.io.File.createTempFile("module", "obj")
+    val absolutePath = tmpFile.getAbsolutePath
+    val module = Sequential()
+
+    module.add(SpatialConvolution(1, 6, 5, 5).setName("conv1"))
+    module.add(Tanh())
+    module.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 2 : filter bank -> squashing -> max pooling
+    module.add(SpatialConvolution(6, 12, 5, 5).setName("conv2"))
+    module.add(Tanh())
+    module.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 3 : standard 2-layer neural network
+    module.add(Reshape(Array(12 * 5 * 5)))
+    module.add(Linear(12 * 5 * 5, 100).setName("l1"))
+    module.add(Tanh())
+    module.add(Linear(100, 6).setName("l2"))
+    module.add(LogSoftMax())
+
+    val module2 = Sequential()
+
+    module2.add(SpatialConvolution(1, 6, 5, 5).setName("conv1"))
+    module2.add(Tanh())
+    module2.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 2 : filter bank -> squashing -> max pooling
+    module2.add(SpatialConvolution(6, 12, 5, 5).setName("conv2"))
+    module2.add(Tanh())
+    module2.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 3 : standard 2-layer neural network
+    module2.add(Reshape(Array(12 * 5 * 5)))
+    module2.add(Linear(12 * 5 * 5, 100).setName("l1"))
+    module2.add(Tanh())
+    module2.add(Linear(100, 6).setName("l2"))
+    module2.add(LogSoftMax())
+
+    module.saveWeights(absolutePath, true)
+
+    module2.loadWeights(absolutePath)
+
+    module.parameters()._1 should be(module2.parameters()._1)
+  }
+
+  "weights save and load with different model definition" should "work properly" in {
+    val tmpFile = java.io.File.createTempFile("module", "obj")
+    val absolutePath = tmpFile.getAbsolutePath
+    val module = Sequential()
+
+    module.add(SpatialConvolution(1, 6, 5, 5).setName("conv1"))
+    module.add(Tanh())
+    module.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 2 : filter bank -> squashing -> max pooling
+    module.add(SpatialConvolution(6, 12, 5, 5).setName("conv2"))
+    module.add(Tanh())
+    module.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 3 : standard 2-layer neural network
+    module.add(Reshape(Array(12 * 5 * 5)))
+    module.add(Linear(12 * 5 * 5, 100).setName("l1"))
+    module.add(Tanh())
+    module.add(Linear(100, 6).setName("l2"))
+    module.add(LogSoftMax())
+
+    val module2 = Sequential()
+
+    module2.add(SpatialConvolution(1, 6, 5, 5).setName("conv1"))
+    module2.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 2 : filter bank -> squashing -> max pooling
+    module2.add(SpatialConvolution(6, 12, 5, 5).setName("conv2"))
+    module2.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 3 : standard 2-layer neural network
+    module2.add(Reshape(Array(12 * 5 * 5)))
+    module2.add(Linear(12 * 5 * 5, 100).setName("l1"))
+    module2.add(Linear(100, 6).setName("l2"))
+
+    module.saveWeights(absolutePath, true)
+
+    module2.loadWeights(absolutePath)
+
+    module.parameters()._1 should be(module2.parameters()._1)
+  }
+
+  "weights save and load with only weight or bias" should "work properly" in {
+    val tmpFile = java.io.File.createTempFile("module", "obj")
+    val absolutePath = tmpFile.getAbsolutePath
+    val module = Sequential()
+
+    module.add(CMul(Array(1, 4, 1, 1)).setName("cmul"))
+    module.add(CAdd(Array(1, 4, 1, 1)).setName("cadd"))
+
+    val module2 = Sequential()
+
+    module2.add(CMul(Array(1, 4, 1, 1)).setName("cmul"))
+    module2.add(CAdd(Array(1, 4, 1, 1)).setName("cadd"))
+
+    module.saveWeights(absolutePath, true)
+
+    module2.loadWeights(absolutePath)
+
+    module.parameters()._1 should be(module2.parameters()._1)
+  }
+
+  "loadModelWeights" should "work properly" in {
+    val module = Sequential()
+
+    module.add(SpatialConvolution(1, 6, 5, 5).setName("conv1"))
+    module.add(Tanh())
+    module.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 2 : filter bank -> squashing -> max pooling
+    module.add(SpatialConvolution(6, 12, 5, 5).setName("conv2"))
+    module.add(Tanh())
+    module.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 3 : standard 2-layer neural network
+    module.add(Reshape(Array(12 * 5 * 5)))
+    module.add(Linear(12 * 5 * 5, 100).setName("l1"))
+    module.add(Tanh())
+    module.add(Linear(100, 6).setName("l2"))
+    module.add(LogSoftMax())
+
+    val module2 = Sequential()
+
+    module2.add(SpatialConvolution(1, 6, 5, 5).setName("conv1"))
+    module2.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 2 : filter bank -> squashing -> max pooling
+    module2.add(SpatialConvolution(6, 12, 5, 5).setName("conv2"))
+    module2.add(SpatialMaxPooling(2, 2, 2, 2))
+    // stage 3 : standard 2-layer neural network
+    module2.add(Reshape(Array(12 * 5 * 5)))
+    module2.add(Linear(12 * 5 * 5, 100).setName("l1"))
+    module2.add(Linear(100, 6).setName("l2"))
+    module.loadModelWeights(module2)
+
+    module.parameters()._1 should be(module2.parameters()._1)
+  }
 }


Code review comments:
diff_hunk: @@ -432,6 +432,75 @@ abstract class AbstractModule[A <: Activity: ClassTag, B <: Activity: ClassTag,
   }
 
   /**
+   * save weights and bias to file
+   * @param path file to save
+   * @param overWrite whether to overwrite or not
+   */
+  def saveWeights(path: String, overWrite: Boolean): Unit = {
+    val parameterTable = getParametersTable()
+    val weightsBiasTable = T()
+    parameterTable.foreach {
+      case (name: String, params: Table) =>
+        val wb = T()
+        if (params.contains("weight")) {
+          wb("weight") = params("weight")
+        }
+        if (params.contains("bias")) {
+          wb("bias") = params("bias")
+        }
+        weightsBiasTable(name) = wb
+      case _ => throw new UnsupportedOperationException("invalid parameter table")
+    }
+    weightsBiasTable.save(path, overWrite)
+  }
+
+  /**
+   * load pretrained weights and bias to current module
+   * @param weightPath file to store weights and bias
+   * @param matchAll whether to match all layers' weights and bias,
+   *                 if not, only load existing pretrained weights and bias
+   * @return current module
+   */
+  def loadWeights(weightPath: String, matchAll: Boolean = true): this.type = {
+    val srcParameter = File.load[Table](weightPath)
+    val targetParameter = getParametersTable()
+    copyWeights(targetParameter, srcParameter)
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala
text: @user1:
perhaps here miss param "matchAll"

@author:
yes

diff_hunk: @@ -432,6 +432,75 @@ abstract class AbstractModule[A <: Activity: ClassTag, B <: Activity: ClassTag,
   }
 
   /**
+   * save weights and bias to file
+   * @param path file to save
+   * @param overWrite whether to overwrite or not
+   */
+  def saveWeights(path: String, overWrite: Boolean): Unit = {
+    val parameterTable = getParametersTable()
+    val weightsBiasTable = T()
+    parameterTable.foreach {
+      case (name: String, params: Table) =>
+        val wb = T()
+        if (params.contains("weight")) {
+          wb("weight") = params("weight")
+        }
+        if (params.contains("bias")) {
+          wb("bias") = params("bias")
+        }
+        weightsBiasTable(name) = wb
+      case _ => throw new UnsupportedOperationException("invalid parameter table")
+    }
+    weightsBiasTable.save(path, overWrite)
+  }
+
+  /**
+   * load pretrained weights and bias to current module
+   * @param weightPath file to store weights and bias
+   * @param matchAll whether to match all layers' weights and bias,
+   *                 if not, only load existing pretrained weights and bias
+   * @return current module
+   */
+  def loadWeights(weightPath: String, matchAll: Boolean = true): this.type = {
+    val srcParameter = File.load[Table](weightPath)
+    val targetParameter = getParametersTable()
+    copyWeights(targetParameter, srcParameter)
+    this
+  }
+
+  /**
+   * copy weights from another model, mapping by layer name
+   * @param srcModel model to copy from
+   * @param matchAll whether to match all layers' weights and bias,
+   * @return current module
+   */
+  def loadModelWeights(srcModel: Module[Float], matchAll: Boolean = true): this.type = {
+    val srcParameters = srcModel.getParametersTable()
+    val parameterTable = getParametersTable()
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala
text: @user1:
similar to "loadweights" function, name "targetParameter" is better