Override copy function in Functional and Sequential models

Kotlin · Dec 19, 2022 · bbf1be7 · bbf1be7
1 parent 90b046d
commit bbf1be7
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/Functional.kt b/tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/Functional.kt
@@ -108,10 +108,10 @@ public class Functional(vararg layers: Layer) : GraphTrainableModel(*layers) {
         file.writeBytes(kGraph.tfGraph.toGraphDef())
     }
 
-    /** Returns a copy of this model. */
-    public fun copy(saveOptimizerState: Boolean = false, copyWeights: Boolean = true): Functional {
+    override fun copy(copiedModelName: String?, saveOptimizerState: Boolean, copyWeights: Boolean): Functional {
         val serializedModel = serializeModel(true)
         val deserializedModel = deserializeFunctionalModel(serializedModel)
+        if (copiedModelName != null) deserializedModel.name = copiedModelName
         if (!copyWeights) return deserializedModel
 
         // TODO: make deep copies, not just links

diff --git a/tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/Sequential.kt b/tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/Sequential.kt
@@ -41,10 +41,10 @@ public class Sequential(vararg layers: Layer) : GraphTrainableModel(*layers) {
         return input to output
     }
 
-    /** Returns a copy of this model. */
-    public fun copy(saveOptimizerState: Boolean = false, copyWeights: Boolean = true): Sequential {
+    override fun copy(copiedModelName: String?, saveOptimizerState: Boolean, copyWeights: Boolean): Sequential {
         val serializedModel = serializeModel(true)
         val deserializedModel = deserializeSequentialModel(serializedModel)
+        if (copiedModelName != null) deserializedModel.name = copiedModelName
         if (!copyWeights) return deserializedModel
 
         deserializedModel.compile(

diff --git a/...orflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/inference/TensorFlowInferenceModel.kt b/...orflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/inference/TensorFlowInferenceModel.kt
@@ -150,7 +150,8 @@ public open class TensorFlowInferenceModel : InferenceModel {
         return copy(copiedModelName = null, saveOptimizerState = false, copyWeights = true)
     }
 
-    public fun copy(
+    /** Returns a copy of this model. */
+    public open fun copy(
         copiedModelName: String? = null,
         saveOptimizerState: Boolean = false, // TODO, check this case
         copyWeights: Boolean = true