fix pytorch featureset, total train sample is incorrect. (intel-analy…

…tics#2421) * fix pytorch featureset * fix style check
dding3 · Jul 2, 2020 · 19aadb0 · 19aadb0
1 parent f795ec0
commit 19aadb0
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 22 deletions.
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/dllib/feature/FeatureSet.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/dllib/feature/FeatureSet.scala
@@ -415,7 +415,7 @@ object PythonFeatureSet{
 class PythonFeatureSet[T: ClassTag](
     dataset: Array[Byte],
     getLoader: (Int, Int, String) => String,
-    getIterator: (String, String) => String,
+    getIterator: (String, String, Boolean) => String,
     getNext: (String) => String,
     inputName: String,
     targetName: String = "",
@@ -441,7 +441,7 @@ class PythonFeatureSet[T: ClassTag](
       cachedRdd.mapPartitions{dataIter =>
         val localLoaderName = getLocalLoader(loaderName)
         val localIterName = getLocalIter(localLoaderName, train)
-        val getIteratorCode = getIterator(localIterName, localLoaderName)
+        val getIteratorCode = getIterator(localIterName, localLoaderName, train)
 
         val nextCode = getNext(localIterName)
         new Iterator[T] {
@@ -478,7 +478,7 @@ class PythonFeatureSet[T: ClassTag](
       cachedRdd.mapPartitions{ dataIter =>
         val localLoaderName = getLocalLoader(loaderName)
         val localIterName = getLocalIter(localLoaderName, train)
-        PythonInterpreter.exec(getIterator(localIterName, localLoaderName))
+        PythonInterpreter.exec(getIterator(localIterName, localLoaderName, train))
         new Iterator[T] {
           val nextCode = getNext(localIterName)
           var alreadyNext = false
@@ -639,7 +639,7 @@ object FeatureSet {
   private[zoo] def python[T: ClassTag](
       dataset: Array[Byte],
       getLoader: (Int, Int, String) => String,
-      getIterator: (String, String) => String,
+      getIterator: (String, String, Boolean) => String,
       getNext: (String) => String,
       inputName: String,
       targetName: String,

diff --git a/...k/dl/src/main/scala/com/intel/analytics/bigdl/dllib/feature/python/PythonFeatureSet.scala b/...k/dl/src/main/scala/com/intel/analytics/bigdl/dllib/feature/python/PythonFeatureSet.scala
@@ -91,7 +91,7 @@ class PythonFeatureSet[T: ClassTag](implicit ev: TensorNumeric[T]) extends Pytho
         |from zoo.util.nest import flatten
         |sess = tf.Session()
         |""".stripMargin
-    def getIterator(iterName: String, loaderName: String): String = {
+    def getIterator(iterName: String, loaderName: String, train: Boolean): String = {
       s"""
          |${iterName} = ${loaderName}.make_one_shot_iterator()
          |""".stripMargin
@@ -116,22 +116,28 @@ class PythonFeatureSet[T: ClassTag](implicit ev: TensorNumeric[T]) extends Pytho
 
   def createFeatureSetFromPyTorch(
         dataloader: Array[Byte]): FeatureSet[MiniBatch[Float]] = {
+    val trainPostfix = "_train"
+    val evalPostfix = "_eval"
     val imports = s"""
                      |from zoo.util.nest import ptensor_to_numpy
                      |import torch
                      |from torch.utils.data import DataLoader
                      |
                      |""".stripMargin
 
-    def getIterator(iterName: String, loaderName: String): String = {
-      s"""
-         |if '${loaderName}_epoch' not in dir():
-         |  ${loaderName}_epoch = 0
-         |else:
-         |  ${loaderName}_epoch += 1
-         |${loaderName}_sampler.set_epoch(${loaderName}_epoch)
-         |${iterName} = enumerate(${loaderName})
-         |""".stripMargin
+    def getIterator(iterName: String, loaderName: String, train: Boolean): String = {
+      if (train) {
+        s"""
+           |if '${loaderName}_epoch' not in dir():
+           |  ${loaderName}_epoch = 0
+           |else:
+           |  ${loaderName}_epoch += 1
+           |${loaderName}_random_sampler.set_epoch(${loaderName}_epoch)
+           |${iterName} = enumerate(${loaderName}${trainPostfix})
+           |""".stripMargin
+      } else {
+        s"${iterName} = enumerate(${loaderName}${evalPostfix})"
+      }
     }
 
     def getNext(iterName: String): String = {
@@ -155,12 +161,10 @@ class PythonFeatureSet[T: ClassTag](implicit ev: TensorNumeric[T]) extends Pytho
            |from torch.utils.data import DataLoader
            |import math
            |
-           |if isinstance(${localLoaderName}.sampler, RandomSampler):
-           |    ${localLoaderName}_sampler=DistributedSampler(${localLoaderName}.dataset,
-           |                                                  ${nodeNumber}, ${partId}, True)
-           |else:
-           |    ${localLoaderName}_sampler=DistributedSequentialSampler(${localLoaderName}.dataset,
-           |                                                  ${nodeNumber}, ${partId})
+           |${localLoaderName}_rand_sampler=DistributedSampler(${localLoaderName}.dataset,
+           |                                              ${nodeNumber}, ${partId}, True)
+           |${localLoaderName}_seq_sampler=DistributedSequentialSampler(${localLoaderName}.dataset,
+           |                                              ${nodeNumber}, ${partId})
            |
            |bs_node = int(math.ceil(${localLoaderName}.batch_size / ${nodeNumber}))
            |
@@ -173,9 +177,11 @@ class PythonFeatureSet[T: ClassTag](implicit ev: TensorNumeric[T]) extends Pytho
            |                "drop_last": ${localLoaderName}.drop_last,
            |                "timeout": ${localLoaderName}.timeout,
            |                "worker_init_fn": ${localLoaderName}.worker_init_fn,
-           |                "sampler": ${localLoaderName}_sampler
+           |                "sampler": ${localLoaderName}_rand_sampler
            |            }
-           |${localLoaderName} = DataLoader(**data_loader_args)
+           |${localLoaderName}${trainPostfix} = DataLoader(**data_loader_args)
+           |data_loader_args["sampler"] = ${localLoaderName}_seq_sampler
+           |${localLoaderName}${evalPostfix} = DataLoader(**data_loader_args)
            |""".stripMargin
     }