Cleanup model copying

api/src/main/kotlin/org/jetbrains/kotlinx/dl/api/inference/InferenceModel.kt

@@ -37,14 +37,9 @@ public interface InferenceModel : AutoCloseable {
     public fun reshape(vararg dims: Long)
 
     /**
-     * Creates a copy.
+     * Creates a copy of this model.
      *
-     * @param [copiedModelName] Set up this name to make a copy with a new name.
      * @return A copied inference model.
      */
-    public fun copy(
-        copiedModelName: String? = null,
-        saveOptimizerState: Boolean = false,
-        copyWeights: Boolean = true
-    ): InferenceModel
+    public fun copy(): InferenceModel
 }

onnx/src/commonMain/kotlin/org/jetbrains/kotlinx/dl/onnx/inference/OnnxInferenceModel.kt

@@ -271,13 +271,8 @@ public open class OnnxInferenceModel private constructor(
         }
     }
 
-    override fun copy(
-        copiedModelName: String?,
-        saveOptimizerState: Boolean,
-        copyWeights: Boolean
-    ): OnnxInferenceModel {
+    override fun copy(): OnnxInferenceModel {
         val model = OnnxInferenceModel(modelSource)
-        model.name = copiedModelName
         if (inputShape != null) {
             model.reshape(*inputDimensions)
         }

tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/Functional.kt

@@ -9,7 +9,6 @@ import org.jetbrains.kotlinx.dl.api.core.layer.Layer
 import org.jetbrains.kotlinx.dl.api.core.layer.core.Input
 import org.jetbrains.kotlinx.dl.api.core.layer.freeze
 import org.jetbrains.kotlinx.dl.api.core.layer.setOutputShape
-import org.jetbrains.kotlinx.dl.api.core.layer.weights
 import org.jetbrains.kotlinx.dl.api.core.util.sortTopologically
 import org.jetbrains.kotlinx.dl.api.inference.keras.*
 import org.tensorflow.Operand
@@ -43,28 +42,26 @@ public class Functional(vararg layers: Layer) : GraphTrainableModel(*layers) {
         return input to output[layers.last()]!!
     }
 
-    /** Returns a copy of this model. */
-    // TODO: support saveOptimizerState=true with assignment of intermediate optimizer state
-    public fun copy(saveOptimizerState: Boolean = false, copyWeights: Boolean = true): Functional {
-        val serializedModel = serializeModel(true)
-        val deserializedModel = deserializeFunctionalModel(serializedModel)
-        if (!copyWeights) {
-            return deserializedModel
-        } else {
-            // TODO: make deep copies, not just links
-            deserializedModel.compile(
-                optimizer = this.optimizer,
-                loss = this.loss,
-                metrics = this.metrics
-            )
-
-            deserializedModel.layers.forEach {
-                it.weights = this.getLayer(it.name).weights
-            }
-
-            deserializedModel.isModelInitialized = true
+    override fun copy(): Functional {
+        return copy(copiedModelName = null, copyOptimizerState = false, copyWeights = true)
+    }
 
-            return deserializedModel
+    /**
+     * Creates a copy of this model.
+     *
+     * @param [copiedModelName] a name for the copy
+     * @param [copyOptimizerState] whether optimizer state needs to be copied
+     * @param [copyWeights] whether model weights need to be copied
+     * @return A copied inference model.
+     */
+    public fun copy(copiedModelName: String? = null,
+                    copyOptimizerState: Boolean = false,
+                    copyWeights: Boolean = true
+    ): Functional {
+        val serializedModel = serializeModel(true)
+        return deserializeFunctionalModel(serializedModel).also { modelCopy ->
+            if (copiedModelName != null) modelCopy.name = copiedModelName
+            if (copyWeights) copyWeightsTo(modelCopy, copyOptimizerState)
         }
     }
 

tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/GraphTrainableModel.kt

@@ -961,6 +961,27 @@ public abstract class GraphTrainableModel(vararg layers: Layer) : TrainableModel
         variable.initializerOperation.run(session)
     }
 
+    protected fun copyWeightsTo(model: GraphTrainableModel, copyOptimizerState: Boolean) {
+        // TODO: make deep copies, not just links
+        model.compile(
+            optimizer = this.optimizer,
+            loss = this.loss,
+            metrics = this.metrics
+        )
+
+        model.layers.forEach {
+            it.weights = this.getLayer(it.name).weights
+        }
+
+        if (copyOptimizerState) {
+            val optimizerVariables = kGraph.variableNames().filter(::isOptimizerVariable)
+            copyVariablesToModel(model, optimizerVariables)
+            model.isOptimizerVariableInitialized = true
+        }
+
+        model.isModelInitialized = true
+    }
+
     /**
      * Return layer by [layerName].
      *

tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/Sequential.kt

@@ -8,7 +8,6 @@ package org.jetbrains.kotlinx.dl.api.core
 import org.jetbrains.kotlinx.dl.api.core.layer.Layer
 import org.jetbrains.kotlinx.dl.api.core.layer.core.Input
 import org.jetbrains.kotlinx.dl.api.core.layer.setOutputShape
-import org.jetbrains.kotlinx.dl.api.core.layer.weights
 import org.jetbrains.kotlinx.dl.api.inference.keras.*
 import org.tensorflow.Operand
 import org.tensorflow.op.core.Placeholder
@@ -41,27 +40,26 @@ public class Sequential(vararg layers: Layer) : GraphTrainableModel(*layers) {
         return input to output
     }
 
-    /** Returns a copy of this model. */
-    // TODO: implement the saving of optimizer state
-    public fun copy(saveOptimizerState: Boolean = false, copyWeights: Boolean = true): Sequential {
-        val serializedModel = serializeModel(true)
-        val deserializedModel = deserializeSequentialModel(serializedModel)
-        if (!copyWeights) {
-            return deserializedModel
-        } else {
-            deserializedModel.compile(
-                optimizer = this.optimizer,
-                loss = this.loss,
-                metrics = this.metrics
-            )
-
-            deserializedModel.layers.forEach {
-                it.weights = this.getLayer(it.name).weights
-            }
-
-            deserializedModel.isModelInitialized = true
+    override fun copy(): Sequential {
+        return copy(copiedModelName = null, copyOptimizerState = false, copyWeights = true)
+    }
 
-            return deserializedModel
+    /**
+     * Creates a copy of this model.
+     *
+     * @param [copiedModelName] a name for the copy
+     * @param [copyOptimizerState] whether optimizer state needs to be copied
+     * @param [copyWeights] whether model weights need to be copied
+     * @return A copied inference model.
+     */
+    public fun copy(copiedModelName: String? = null,
+                    copyOptimizerState: Boolean = false,
+                    copyWeights: Boolean = true
+    ): Sequential {
+        val serializedModel = serializeModel(true)
+        return deserializeSequentialModel(serializedModel).also { modelCopy ->
+            if (copiedModelName != null) modelCopy.name = copiedModelName
+            if (copyWeights) copyWeightsTo(modelCopy, copyOptimizerState)
         }
     }
 

tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/inference/TensorFlowInferenceModel.kt

@@ -145,11 +145,12 @@ public open class TensorFlowInferenceModel : InferenceModel {
         }
     }
 
-    override fun copy(
-        copiedModelName: String?,
-        saveOptimizerState: Boolean, // TODO, check this case
-        copyWeights: Boolean
-    ): TensorFlowInferenceModel {
+    override fun copy(): TensorFlowInferenceModel {
+        return copy(copiedModelName = null)
+    }
+
+    /** Returns a copy of this model. */
+    public fun copy(copiedModelName: String? = null): TensorFlowInferenceModel {
         val model = TensorFlowInferenceModel()
         model.kGraph = this.kGraph.copy()
         model.tf = Ops.create(model.kGraph.tfGraph)
@@ -158,27 +159,21 @@ public open class TensorFlowInferenceModel : InferenceModel {
         model.input = input
         model.output = output
         if (copiedModelName != null) model.name = name
-        // TODO: check that tensors are closed after usage
-        if (copyWeights) {
-            val modelWeightsExtractorRunner = session.runner()
-            val variableNames = kGraph.variableNames()
-            check(variableNames.isNotEmpty()) {
-                "Found 0 variable names in TensorFlow graph $kGraph. " +
-                        "If copied model has no weights, set flag `copyWeights` to `false`."
-            }
+        copyVariablesToModel(model, kGraph.variableNames())
+        model.isModelInitialized = true
+        return model
+    }
 
-            val variableNamesToCopy = variableNames.filter { variableName ->
-                saveOptimizerState || !isOptimizerVariable(variableName)
-            }
-            variableNamesToCopy.forEach(modelWeightsExtractorRunner::fetch)
-            val modelWeights = variableNamesToCopy.zip(modelWeightsExtractorRunner.run()).toMap()
+    protected fun copyVariablesToModel(model: TensorFlowInferenceModel, variableNames: List<String>) {
+        if (variableNames.isEmpty()) return
 
-            model.loadVariables(modelWeights.keys) { variableName, _ ->
-                modelWeights[variableName]!!.use { it.convertTensorToMultiDimArray() }
-            }
+        val modelWeightsExtractorRunner = session.runner()
+        variableNames.forEach(modelWeightsExtractorRunner::fetch)
+        val modelWeights = variableNames.zip(modelWeightsExtractorRunner.run()).toMap()
+
+        model.loadVariables(modelWeights.keys) { variableName, _ ->
+            modelWeights[variableName]!!.use { it.convertTensorToMultiDimArray() }
         }
-        model.isModelInitialized = true
-        return model
     }
 
     /**

tensorflow/src/test/kotlin/org/jetbrains/kotlinx/dl/api/core/integration/TransferLearningFromKerasTest.kt

@@ -306,7 +306,7 @@ class TransferLearningTest : IntegrationTest() {
 
             it.loadWeights(hdfFile)
 
-            val copy = it.copy()
+            val copy = it.copy(copyOptimizerState = false, copyWeights = true)
             assertTrue(copy.layers.size == 11)
             copy.close()