save TorchNet as Pytorch script module (intel-analytics#1564)

* support save pytorch model to script * unit test * use temp folder * add import * correct evaluate * import * style fix
dding3 · Sep 25, 2019 · 07f4ebe · 07f4ebe
1 parent 13ea6f8
commit 07f4ebe
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/bigdl/dllib/src/main/scala/com/intel/analytics/bigdl/dllib/inference/net/TorchNet.scala b/bigdl/dllib/src/main/scala/com/intel/analytics/bigdl/dllib/inference/net/TorchNet.scala
@@ -69,6 +69,7 @@ class TorchNet private(private val modelHolder: TorchModelHolder)
   }
 
   override def evaluate(): this.type = {
+    nativeRef
     super.evaluate()
     if (!weights.isEmpty) {
       PytorchModel.updateWeightNative(nativeRef, weights.storage().array())
@@ -144,6 +145,14 @@ class TorchNet private(private val modelHolder: TorchModelHolder)
     super.finalize()
     PytorchModel.releaseModelNative(nativeRef)
   }
+
+  /**
+   * export the model to path as a torch script module.
+   */
+  def savePytorch(path : String, overWrite: Boolean = false): Unit = {
+    PytorchModel.updateWeightNative(this.nativeRef, weights.storage().array())
+    PytorchModel.saveModelNative(nativeRef, path)
+  }
 }
 
 object TorchNet {

diff --git a/...ib/src/main/scala/com/intel/analytics/bigdl/dllib/inference/net/python/PythonZooNet.scala b/...ib/src/main/scala/com/intel/analytics/bigdl/dllib/inference/net/python/PythonZooNet.scala
@@ -301,4 +301,8 @@ class PythonZooNet[T: ClassTag](implicit ev: TensorNumeric[T]) extends PythonZoo
     TorchCriterion(lossPath)
   }
 
+  def torchNetSavePytorch(torchnet: TorchNet, path: String): Unit = {
+    torchnet.savePytorch(path)
+  }
+
 }