Add big TF model support (intel-analytics#2974)

* test * add print * add graphdef print * update * update * test * update * add print * update broadcast * fix spark file location * update broadcast * fix extra init * update broadcast * update property * update broadcast * restore broadcast * restore clone * fix clone * fix get extra * update collect weights * update * update * update property * update get extra param * update * update * restore * remove unused import * fix style * add methods * fix style
Le-Zheng · Oct 23, 2020 · 08688e8 · 08688e8
1 parent 885a4d7
commit 08688e8
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/scala/orca/src/main/scala/com/intel/analytics/bigdl/orca/net/TFNet.scala b/scala/orca/src/main/scala/com/intel/analytics/bigdl/orca/net/TFNet.scala
@@ -62,6 +62,9 @@ class TFNet(private val graphDef: TFGraphHolder,
   implicit val ev = TensorNumeric.NumericFloat
   implicit val tag: ClassTag[Float] = ClassTag.Float
 
+  System.setProperty("bigdl.ModelBroadcastFactory",
+    "com.intel.analytics.zoo.tfpark.TFModelBroadcastFactory")
+
   @transient
   private lazy val tensorManager = new TFResourceManager()
 

diff --git a/scala/orca/src/main/scala/com/intel/analytics/bigdl/orca/net/TFNetForInference.scala b/scala/orca/src/main/scala/com/intel/analytics/bigdl/orca/net/TFNetForInference.scala
@@ -49,6 +49,9 @@ private[zoo] class TFNetForInference(graphRunner: GraphRunner,
   implicit val ev = TensorNumeric.NumericFloat
   implicit val tag: ClassTag[Float] = ClassTag.Float
 
+  System.setProperty("bigdl.ModelBroadcastFactory",
+    "com.intel.analytics.zoo.tfpark.TFModelBroadcastFactory")
+
   override def parameters(): (Array[Tensor[Float]], Array[Tensor[Float]]) = {
     (weights, gradWeights)
   }